{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ef4630a-9860-4de2-945e-3fbc71650cbb",
   "metadata": {},
   "source": [
    "# 🚀 Unity Catalog AI: Advanced Developer Workflow Tutorial\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Unity Catalog AI 0.2.0 and 0.3.0 introduce groundbreaking enhancements to the developer experience, transforming how data scientists and ML engineers create, test, and deploy functions for AI applications. This tutorial explores these powerful workflow improvements with a focus on scientific computing applications.\n",
    "\n",
    "Whether you're building AI agents that leverage domain-specific knowledge or creating reusable function libraries for data science workflows, the tools introduced in these releases will dramatically improve your productivity and code quality.\n",
    "\n",
    "## 🎯 What You'll Learn\n",
    "\n",
    "This tutorial guides you through a complete developer workflow using real-world astrophysics examples:\n",
    "\n",
    "- **Function Creation & Registration**: Build sophisticated scientific functions with rich documentation and parameter validation\n",
    "- **Function Introspection**: Leverage the new ability to retrieve, examine, and modify existing function definitions\n",
    "- **Execution & Testing**: Validate functions with comprehensive testing across parameter spaces\n",
    "- **Function Wrapping**: Minimize code duplication by consolidating multiple specialized functions\n",
    "- **Sandbox Protection**: Ensure system stability with process isolation and resource limits\n",
    "\n",
    "## 🔍 Why This Matters\n",
    "\n",
    "The enhancements in Unity Catalog AI 0.2.0 and 0.3.0 address critical pain points in the AI function development lifecycle:\n",
    "\n",
    "- **Reduced Code Duplication**: Create functions once and reuse them across multiple contexts\n",
    "- **Improved Safety**: Protect your systems from resource exhaustion and unintended behavior\n",
    "- **Enhanced Maintainability**: Update existing functions without rewriting them from scratch\n",
    "- **Better Integration**: Build functions that work seamlessly with modern AI frameworks\n",
    "\n",
    "Let's explore how these powerful features can transform your workflow for creating AI-powered scientific applications!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a28e1928-43e3-4157-8c33-66fc2ce5ecf6",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "\n",
    "### Unity Catalog Server\n",
    "\n",
    "To run this tutorial, you will need:\n",
    "\n",
    "1. A clone of the [Unity Catalog repository](https://github.com/unitycatalog/unitycatalog).\n",
    "\n",
    "    ```sh\n",
    "    git clone https://github.com/unitycatalog/unitycatalog\n",
    "    ```\n",
    "\n",
    "2. JDK-17 installed on your system (in order to build and run the Unity Catalog services)<sup>1</sup>\n",
    "3. [Docker Desktop](https://www.docker.com/products/docker-desktop/) (recommended; you can also install the docker engine yourself)\n",
    "4. At the repository root, run:\n",
    "\n",
    "    ```sh\n",
    "    docker compose up\n",
    "    ```\n",
    "\n",
    "<sup>1</sup> (New to managing Java environments? [jenv](https://github.com/jenv/jenv) is a fantastic tool that can help!) \n",
    "\n",
    "\n",
    "### Python Requirements\n",
    "\n",
    "Unity Catalog AI requires Python 3.10+. \n",
    "\n",
    "```sh\n",
    "pip install -U unitycatalog-ai\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d931052c-a35a-479a-be71-66476b01af50",
   "metadata": {},
   "source": [
    "## Setting Up Your Unity Catalog Environment\n",
    "\n",
    "Before we dive into Unity Catalog AI's powerful features, we need to establish a connection to our Unity Catalog server. \n",
    "\n",
    "The 0.3.0 release introduces significant security and execution enhancements. Notice how we're initializing our client with `execution_mode=\"sandbox\"` - this is a new security feature that executes functions in an isolated process pool rather than in the main process. This provides:\n",
    "\n",
    "- Better resource management with CPU limitations\n",
    "- Timeout restrictions to prevent runaway processes\n",
    "- Blocking of modules that might access sensitive system resources\n",
    "\n",
    "We'll also create a catalog and schema to organize our astrophysics functions, demonstrating the utility methods added in recent releases that simplify catalog and schema management."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d52d94bd-b664-4606-bfb7-06212f8ef177",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SchemaInfo(name='physics_functions', catalog_name='astro_catalog', comment='Schema for astrophysics examples', properties={}, full_name='astro_catalog.physics_functions', owner=None, created_at=1742840579002, created_by=None, updated_at=1742840579002, updated_by=None, schema_id='26d02587-86d3-45ed-a8ce-8616cfb02517')"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from unitycatalog.ai.core.client import UnitycatalogFunctionClient\n",
    "from unitycatalog.client import ApiClient, Configuration\n",
    "\n",
    "# Configure the Unity Catalog OSS client\n",
    "config = Configuration(host=\"http://localhost:8080/api/2.1/unity-catalog\")\n",
    "api_client = ApiClient(configuration=config)\n",
    "\n",
    "client = UnitycatalogFunctionClient(api_client=api_client, execution_mode=\"sandbox\")\n",
    "\n",
    "CATALOG = \"astro_catalog\"\n",
    "SCHEMA = \"physics_functions\"\n",
    "\n",
    "# Create catalog and schema if they don't already exist\n",
    "client.uc.create_catalog(name=CATALOG, comment=\"Catalog for astrophysics calculations\")\n",
    "client.uc.create_schema(\n",
    "    name=SCHEMA, catalog_name=CATALOG, comment=\"Schema for astrophysics examples\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9936c6d8-3b21-4787-8709-6989f93fe9c4",
   "metadata": {},
   "source": [
    "## Crafting Your First Unity Catalog Function\n",
    "\n",
    "The first step in our development workflow is creating an initial function to register in Unity Catalog. We'll define a function for calculating stellar properties based on astrophysical parameters - a perfect example to showcase how Unity Catalog AI handles complex scientific functions.\n",
    "\n",
    "This function simulates stellar evolution in a simplified way, calculating key properties like:\n",
    "- Current mass (after accounting for stellar mass loss over time)\n",
    "- Luminosity (how bright the star is compared to our Sun)\n",
    "- Surface temperature (in Kelvin)\n",
    "- Evolutionary stage (from protostar through main sequence and beyond to end states like white dwarf or black hole)\n",
    "\n",
    "While this uses simplified models compared to sophisticated stellar evolution codes used in astrophysics research, it demonstrates how domain expertise can be packaged into a callable function that AI assistants can leverage.\n",
    "\n",
    "Notice the comprehensive docstring in our function definition. This is critical for two reasons:\n",
    "\n",
    "1. **GenAI Integration**: Unity Catalog AI parses these docstrings to create proper tool definitions when integrating with LLMs like Claude, Gemini, or GPT models. The better your documentation, the more effectively AI agents can understand and use your function.\n",
    "\n",
    "2. **Developer Experience**: Clear parameter descriptions, return value documentation, and usage examples make your functions more maintainable and easier for other team members to understand.\n",
    "\n",
    "The 0.2.0 and 0.3.0 releases include enhanced docstring parsing abilities, giving you better control over how your functions appear as tools to AI agents. This function includes parameter validation, meaningful error messages, and clear documentation of parameter constraints - all best practices when creating functions for AI systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "692c07c2-1a09-4e85-a984-efb935e6286e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_stellar_properties(initial_mass: float, metallicity: float, age: float) -> dict:\n",
    "    \"\"\"\n",
    "    Calculate key stellar characteristics based on initial mass, metallicity, and age.\n",
    "    This function provides a simplified model of stellar evolution and properties.\n",
    "\n",
    "    Args:\n",
    "        initial_mass (float): Initial stellar mass in solar masses (M☉).\n",
    "            Valid range: 0.1 to 150 solar masses.\n",
    "        metallicity (float): Stellar metallicity (Z), representing the fraction\n",
    "            of stellar mass composed of elements heavier than hydrogen and helium.\n",
    "            Valid range: 0.0001 to 0.04.\n",
    "        age (float): Stellar age in billions of years.\n",
    "            Valid range: 0 to 13.8 (age of the universe).\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary of stellar characteristics including:\n",
    "            - current_mass: Remaining stellar mass in solar masses\n",
    "            - luminosity: Current stellar luminosity relative to the Sun\n",
    "            - surface_temperature: Effective surface temperature in Kelvin\n",
    "            - stellar_stage: Textual description of the star's evolutionary stage\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If input parameters are outside physically meaningful ranges.\n",
    "\n",
    "    Example:\n",
    "        >>> result = calculate_stellar_properties(2.5, 0.02, 1.0)\n",
    "        >>> print(result[\"luminosity\"])\n",
    "    \"\"\"\n",
    "    # Input validation\n",
    "    if not (0.1 <= initial_mass <= 150):\n",
    "        raise ValueError(\"Initial mass must be between 0.1 and 150 solar masses\")\n",
    "\n",
    "    if not (0.0001 <= metallicity <= 0.04):\n",
    "        raise ValueError(\"Metallicity must be between 0.0001 and 0.04\")\n",
    "\n",
    "    if not (0 <= age <= 13.8):\n",
    "        raise ValueError(\"Stellar age must be between 0 and 13.8 billion years\")\n",
    "\n",
    "    # Calculate stellar lifetime (in billions of years)\n",
    "    # Massive stars burn fuel much faster - lifetime scales roughly as M^-2.5\n",
    "    main_sequence_lifetime = 10.0 * (initial_mass / 1.0) ** -2.5\n",
    "\n",
    "    # Adjust lifetime based on metallicity (metal-poor stars live slightly longer)\n",
    "    metallicity_factor = (metallicity / 0.02) ** 0.3\n",
    "    main_sequence_lifetime = main_sequence_lifetime / metallicity_factor\n",
    "\n",
    "    # Calculate relative age (how far through its lifetime the star is)\n",
    "    relative_age = min(age / main_sequence_lifetime, 5.0)  # Cap at 5x lifetime for remnants\n",
    "\n",
    "    # Determine current mass based on evolutionary stage\n",
    "    if relative_age < 0.1:  # Very early evolution\n",
    "        # Minimal mass loss during early stages\n",
    "        current_mass = initial_mass * (1.0 - 0.01 * relative_age)\n",
    "    elif relative_age < 1.0:  # Main sequence\n",
    "        # Gradual mass loss during main sequence\n",
    "        # More massive stars lose more (Wolf-Rayet stars, etc)\n",
    "        main_sequence_loss = 0.05 * (initial_mass / 10.0) ** 1.5\n",
    "        current_mass = initial_mass * (1.0 - main_sequence_loss * relative_age)\n",
    "    elif relative_age < 1.1:  # End of main sequence/subgiant\n",
    "        # Increased mass loss as star expands\n",
    "        current_mass = initial_mass * (0.95 - 0.1 * (relative_age - 1.0))\n",
    "    elif relative_age < 1.3:  # Red giant/AGB phase - dramatic mass loss\n",
    "        # Up to 50-80% mass loss during giant phases\n",
    "        giant_phase_factor = 0.5 if initial_mass < 8.0 else 0.3\n",
    "        current_mass = initial_mass * (0.85 - giant_phase_factor * (relative_age - 1.1) / 0.2)\n",
    "    else:  # Remnant\n",
    "        # Determine final remnant mass based on initial mass\n",
    "        if initial_mass < 8.0:  # White dwarf territory\n",
    "            # Most stars become 0.5-1.0 M☉ white dwarfs\n",
    "            current_mass = 0.5 + 0.1 * initial_mass\n",
    "            current_mass = min(current_mass, 1.4)  # Chandrasekhar limit\n",
    "        elif initial_mass < 20.0:  # Neutron star territory\n",
    "            # Stars 8-20 M☉ tend to make ~1.4-2.0 M☉ neutron stars\n",
    "            current_mass = 1.4 + 0.03 * (initial_mass - 8.0)\n",
    "        else:  # Black hole territory\n",
    "            # Stars above 20 M☉ make black holes roughly 10-20% of initial mass\n",
    "            current_mass = max(3.0, 0.2 * initial_mass)\n",
    "\n",
    "    # Luminosity calculation\n",
    "    if relative_age < 1.0:  # Main sequence luminosity\n",
    "        if current_mass < 0.43:\n",
    "            luminosity = current_mass**2.3\n",
    "        elif current_mass < 2:\n",
    "            luminosity = current_mass**4\n",
    "        elif current_mass < 20:\n",
    "            luminosity = current_mass**3.5\n",
    "        else:\n",
    "            luminosity = current_mass**3.0\n",
    "    elif relative_age < 1.3:  # Giant phase luminosity spike\n",
    "        luminosity = 100 * current_mass**2\n",
    "    else:  # Remnant luminosity\n",
    "        if current_mass < 1.4:  # White dwarf\n",
    "            # White dwarfs cool over time\n",
    "            cooling_factor = max(0.001, 1.0 - (relative_age - 1.3))\n",
    "            luminosity = 0.1 * cooling_factor\n",
    "        elif current_mass < 3.0:  # Neutron star\n",
    "            luminosity = 0.01  # Very low luminosity\n",
    "        else:  # Black hole\n",
    "            # Accretion-dependent, but generally near-zero for isolated black holes\n",
    "            luminosity = 0.0\n",
    "\n",
    "    # Surface temperature estimation\n",
    "    if relative_age < 0.1:  # Protostar\n",
    "        surface_temperature = 3000 + 1000 * relative_age * 10\n",
    "    elif relative_age < 1.0:  # Main sequence\n",
    "        if current_mass < 0.5:\n",
    "            surface_temperature = 3000 + 1000 * current_mass\n",
    "        elif current_mass < 1.5:\n",
    "            surface_temperature = 5000 + 1000 * (current_mass - 0.5)\n",
    "        else:\n",
    "            surface_temperature = 6500 + 5000 * (current_mass / 10.0) ** 0.5\n",
    "    elif relative_age < 1.2:  # Giant phase\n",
    "        surface_temperature = 4000 - 1000 * (relative_age - 1.0) * 5\n",
    "    else:  # Remnant\n",
    "        if current_mass < 1.4:  # White dwarf\n",
    "            surface_temperature = 20000 * max(0.1, 1.0 - 0.1 * (relative_age - 1.3))\n",
    "        elif current_mass < 3.0:  # Neutron star\n",
    "            surface_temperature = 1000000  # Very hot surface\n",
    "        else:  # Black hole\n",
    "            surface_temperature = 0  # No surface\n",
    "\n",
    "    # Stellar stage classification\n",
    "    if relative_age < 0.01:\n",
    "        stellar_stage = \"Protostar\"\n",
    "    elif relative_age < 1.0:\n",
    "        stellar_stage = \"Main Sequence\"\n",
    "    elif relative_age < 1.05:\n",
    "        stellar_stage = \"Subgiant\"\n",
    "    elif relative_age < 1.2:\n",
    "        stellar_stage = \"Red Giant\"\n",
    "    elif relative_age < 1.3:\n",
    "        stellar_stage = \"Asymptotic Giant Branch\"\n",
    "    else:  # Remnant\n",
    "        if current_mass < 1.4:\n",
    "            stellar_stage = \"White Dwarf\"\n",
    "        elif current_mass < 3.0:\n",
    "            stellar_stage = \"Neutron Star\"\n",
    "        else:\n",
    "            stellar_stage = \"Black Hole\"\n",
    "\n",
    "    return {\n",
    "        \"current_mass\": current_mass,\n",
    "        \"luminosity\": luminosity,\n",
    "        \"surface_temperature\": surface_temperature,\n",
    "        \"stellar_stage\": stellar_stage,\n",
    "        \"main_sequence_lifetime\": main_sequence_lifetime,  # In billions of years\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ff1c4c0-bbc2-4e3a-a017-5a6de7d28f35",
   "metadata": {},
   "source": [
    "## Registering Your First Function with Unity Catalog\n",
    "\n",
    "Now that we've created our stellar properties function, it's time to register it with Unity Catalog. This step makes our function accessible to AI assistants, Databricks notebooks, and other applications that connect to the Unity Catalog.\n",
    "\n",
    "The `create_python_function` method handles all the complexity of registering the function. Behind the scenes, it:\n",
    "\n",
    "1. **Analyzes the function signature** to extract parameter types and return values\n",
    "2. **Parses the docstring** to extract descriptions, parameter details, and examples\n",
    "3. **Converts the function** to a format that can be stored and executed by Unity Catalog\n",
    "4. **Registers the function** in the specified catalog and schema\n",
    "\n",
    "The `replace=True` parameter allows us to update the function if it already exists, which is particularly useful during development when we're iterating on function implementations.\n",
    "\n",
    "After registration, we store the fully qualified function name (`{catalog}.{schema}.{function_name}`) for later use. This naming convention is similar to the table naming conventions in databases, providing a hierarchical organization that helps manage complex function libraries.\n",
    "\n",
    "This registration process creates a clean separation between function development and function execution. You can develop and test functions locally, then register them for broader use once they're ready."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9159661c-5a2e-4c3e-8fe6-40df6e88354e",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.create_python_function(\n",
    "    func=calculate_stellar_properties,\n",
    "    catalog=CATALOG,\n",
    "    schema=SCHEMA,\n",
    "    replace=True,\n",
    ")\n",
    "\n",
    "stellar_func_name = f\"{CATALOG}.{SCHEMA}.calculate_stellar_properties\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b3cae68-2331-44ad-9033-a5454f0a41ae",
   "metadata": {},
   "source": [
    "## Registering and Validating Your Function\n",
    "\n",
    "Now that we've defined our stellar evolution function, let's register it with Unity Catalog and validate it with some test cases. But first, an important disclaimer:\n",
    "\n",
    "> **Disclaimer**: The stellar evolution model presented here is a dramatic simplification of actual astrophysical processes. Real stellar evolution codes like MESA, STARLIB, or GENEC solve complex sets of differential equations involving nuclear reaction networks, detailed opacity tables, convection modeling, and more. These codes often contain 100,000+ lines of code developed by research teams over decades. Our model captures only the broad patterns of stellar evolution for demonstration purposes.\n",
    "\n",
    "This section demonstrates one of the most critical aspects of the Unity Catalog AI developer workflow: **testing and validation before deployment**. When working with scientific or analytical functions that will be called by AI agents, it's essential to verify:\n",
    "\n",
    "1. **Parameter boundaries** work as expected\n",
    "2. **Edge cases** produce reasonable results\n",
    "3. **Function outputs** are in expected ranges and formats\n",
    "\n",
    "The 0.3.0 release makes this validation process much easier with the enhanced sandbox execution mode. As you can see from our test scenarios, we're examining:\n",
    "- A Sun-like star (1.0 M☉)\n",
    "- A massive O-type star (25.0 M☉)\n",
    "- A low-mass red dwarf (0.3 M☉)\n",
    "- A very massive star that has evolved to a black hole (90.5 M☉)\n",
    "\n",
    "Verifying these different scenarios helps ensure our function behaves correctly across its parameter space before we expose it to AI systems that might explore edge cases we didn't anticipate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ed6a6084-4879-462c-a043-b883dae09f0a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Stellar Model for 1.0 M☉ star:\n",
      "\tcurrent_mass: 0.9992726761381613\n",
      "\tluminosity: 0.9970938770139077\n",
      "\tsurface_temperature: 5499.272676138161\n",
      "\tstellar_stage: Main Sequence\n",
      "\tmain_sequence_lifetime: 10.0\n",
      "\n",
      "Stellar Model for 25.0 M☉ star:\n",
      "\tcurrent_mass: 22.16718363676904\n",
      "\tluminosity: 10892.600037747503\n",
      "\tsurface_temperature: 13944.323951301596\n",
      "\tstellar_stage: Main Sequence\n",
      "\tmain_sequence_lifetime: 0.0034884427442219237\n",
      "\n",
      "Stellar Model for 0.3 M☉ star:\n",
      "\tcurrent_mass: 0.2999397973798912\n",
      "\tluminosity: 0.06268713412999312\n",
      "\tsurface_temperature: 3200.675400362578\n",
      "\tstellar_stage: Main Sequence\n",
      "\tmain_sequence_lifetime: 498.3171819730835\n",
      "\n",
      "Stellar Model for 90.5 M☉ star:\n",
      "\tcurrent_mass: 18.1\n",
      "\tluminosity: 0.0\n",
      "\tsurface_temperature: 0\n",
      "\tstellar_stage: Black Hole\n",
      "\tmain_sequence_lifetime: 0.00019453446251488867\n"
     ]
    }
   ],
   "source": [
    "import ast\n",
    "\n",
    "scenarios = [\n",
    "    # Sun-like star\n",
    "    (1.0, 0.02, 4.6),\n",
    "    # Massive star\n",
    "    (25.0, 0.015, 0.002),\n",
    "    # Low-mass star\n",
    "    (0.3, 0.001, 10.0),\n",
    "    # Old Massive Star that has evolved to a singularity\n",
    "    (90.5, 0.005, 5.0),\n",
    "]\n",
    "\n",
    "for mass, metallicity, age in scenarios:\n",
    "    result = client.execute_function(\n",
    "        function_name=stellar_func_name,\n",
    "        parameters={\n",
    "            \"initial_mass\": mass,\n",
    "            \"metallicity\": metallicity,\n",
    "            \"age\": age,\n",
    "        },\n",
    "    )\n",
    "    print(f\"\\nStellar Model for {mass} M☉ star:\")\n",
    "    for key, value in ast.literal_eval(result.value).items():\n",
    "        print(f\"\\t{key}: {value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bd13337-9f3e-47a6-8774-a7a99930697f",
   "metadata": {},
   "source": [
    "## Fetching Function Source Code with get_function_source()\n",
    "\n",
    "One of the most powerful additions in the 0.3.0 release is the ability to retrieve the source code of existing Unity Catalog functions. This capability transforms the development workflow by enabling:\n",
    "\n",
    "1. **Code Inspection**: Quickly review how existing functions work without having to search through your codebase\n",
    "2. **Iterative Development**: Make improvements to functions that are already registered\n",
    "3. **Knowledge Transfer**: Learn from functions written by other team members\n",
    "4. **Debugging**: Examine function implementations when unexpected results occur\n",
    "\n",
    "The `get_function_source()` API solves a significant pain point in the Unity Catalog development workflow. Behind the scenes, Unity Catalog stores functions as complex `FunctionInfo` objects that split the function signature, docstring, and body into separate constituent parts within a nested structure. Manually extracting and reconstructing a complete function from this representation is tedious, error-prone, and requires deep knowledge of the internal data model.\n",
    "\n",
    "With `get_function_source()`, you get back a properly formatted Python function as a string - ready to be modified or used as a template for new functions. This seemingly simple capability eliminates what was previously a frustrating and error-prone manual process of copying metadata fields from complex nested objects.\n",
    "\n",
    "In the example above, we retrieve the full definition of our stellar properties function. Notice that the entire function, including comments and docstrings, is preserved exactly as it was registered. This allows us to make targeted improvements while maintaining the function's interface and documentation.\n",
    "\n",
    "In real-world scientific and analytics workflows, functions often evolve as models are refined and improved. The `get_function_source()` API makes this evolution smoother by providing a clear starting point for enhancements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a4f374b3-7635-48fe-bab9-082e3aede1b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def calculate_stellar_properties(initial_mass: float, metallicity: float, age: float) -> dict:\n",
      "    \"\"\"\n",
      "    Calculate key stellar characteristics based on initial mass, metallicity, and age. This function\n",
      "provides a simplified model of stellar evolution and properties.\n",
      "    \n",
      "    Args:\n",
      "        initial_mass: Initial stellar mass in solar masses (M☉). Valid range: 0.1 to 150 solar masses.\n",
      "        metallicity: Stellar metallicity (Z), representing the fraction of stellar mass composed of\n",
      "          elements heavier than hydrogen and helium.  Valid range: 0.0001 to 0.04.\n",
      "        age: Stellar age in billions of years. Valid range: 0 to 13.8 (age of the universe).\n",
      "    \n",
      "    Returns:\n",
      "        dict\n",
      "    \"\"\"\n",
      "    if not (0.1 <= initial_mass <= 150):\n",
      "        raise ValueError(\"Initial mass must be between 0.1 and 150 solar masses\")\n",
      "\n",
      "    if not (0.0001 <= metallicity <= 0.04):\n",
      "        raise ValueError(\"Metallicity must be between 0.0001 and 0.04\")\n",
      "\n",
      "    if not (0 <= age <= 13.8):\n",
      "        raise ValueError(\"Stellar age must be between 0 and 13.8 billion years\")\n",
      "\n",
      "    # Calculate stellar lifetime (in billions of years)\n",
      "    # Massive stars burn fuel much faster - lifetime scales roughly as M^-2.5\n",
      "    main_sequence_lifetime = 10.0 * (initial_mass / 1.0) ** -2.5\n",
      "\n",
      "    # Adjust lifetime based on metallicity (metal-poor stars live slightly longer)\n",
      "    metallicity_factor = (metallicity / 0.02) ** 0.3\n",
      "    main_sequence_lifetime = main_sequence_lifetime / metallicity_factor\n",
      "\n",
      "    # Calculate relative age (how far through its lifetime the star is)\n",
      "    relative_age = min(age / main_sequence_lifetime, 5.0)  # Cap at 5x lifetime for remnants\n",
      "\n",
      "    # Determine current mass based on evolutionary stage\n",
      "    if relative_age < 0.1:  # Very early evolution\n",
      "        # Minimal mass loss during early stages\n",
      "        current_mass = initial_mass * (1.0 - 0.01 * relative_age)\n",
      "    elif relative_age < 1.0:  # Main sequence\n",
      "        # Gradual mass loss during main sequence\n",
      "        # More massive stars lose more (Wolf-Rayet stars, etc)\n",
      "        main_sequence_loss = 0.05 * (initial_mass / 10.0) ** 1.5\n",
      "        current_mass = initial_mass * (1.0 - main_sequence_loss * relative_age)\n",
      "    elif relative_age < 1.1:  # End of main sequence/subgiant\n",
      "        # Increased mass loss as star expands\n",
      "        current_mass = initial_mass * (0.95 - 0.1 * (relative_age - 1.0))\n",
      "    elif relative_age < 1.3:  # Red giant/AGB phase - dramatic mass loss\n",
      "        # Up to 50-80% mass loss during giant phases\n",
      "        giant_phase_factor = 0.5 if initial_mass < 8.0 else 0.3\n",
      "        current_mass = initial_mass * (0.85 - giant_phase_factor * (relative_age - 1.1) / 0.2)\n",
      "    else:  # Remnant\n",
      "        # Determine final remnant mass based on initial mass\n",
      "        if initial_mass < 8.0:  # White dwarf territory\n",
      "            # Most stars become 0.5-1.0 M☉ white dwarfs\n",
      "            current_mass = 0.5 + 0.1 * initial_mass\n",
      "            current_mass = min(current_mass, 1.4)  # Chandrasekhar limit\n",
      "        elif initial_mass < 20.0:  # Neutron star territory\n",
      "            # Stars 8-20 M☉ tend to make ~1.4-2.0 M☉ neutron stars\n",
      "            current_mass = 1.4 + 0.03 * (initial_mass - 8.0)\n",
      "        else:  # Black hole territory\n",
      "            # Stars above 20 M☉ make black holes roughly 10-20% of initial mass\n",
      "            current_mass = max(3.0, 0.2 * initial_mass)\n",
      "\n",
      "    # Luminosity calculation\n",
      "    if relative_age < 1.0:  # Main sequence luminosity\n",
      "        if current_mass < 0.43:\n",
      "            luminosity = current_mass ** 2.3\n",
      "        elif current_mass < 2:\n",
      "            luminosity = current_mass ** 4\n",
      "        elif current_mass < 20:\n",
      "            luminosity = current_mass ** 3.5\n",
      "        else:\n",
      "            luminosity = current_mass ** 3.0\n",
      "    elif relative_age < 1.3:  # Giant phase luminosity spike\n",
      "        luminosity = 100 * current_mass ** 2\n",
      "    else:  # Remnant luminosity\n",
      "        if current_mass < 1.4:  # White dwarf\n",
      "            # White dwarfs cool over time\n",
      "            cooling_factor = max(0.001, 1.0 - (relative_age - 1.3))\n",
      "            luminosity = 0.1 * cooling_factor\n",
      "        elif current_mass < 3.0:  # Neutron star\n",
      "            luminosity = 0.01  # Very low luminosity\n",
      "        else:  # Black hole\n",
      "            # Accretion-dependent, but generally near-zero for isolated black holes\n",
      "            luminosity = 0.0\n",
      "\n",
      "    # Surface temperature estimation\n",
      "    if relative_age < 0.1:  # Protostar\n",
      "        surface_temperature = 3000 + 1000 * relative_age * 10\n",
      "    elif relative_age < 1.0:  # Main sequence\n",
      "        if current_mass < 0.5:\n",
      "            surface_temperature = 3000 + 1000 * current_mass\n",
      "        elif current_mass < 1.5:\n",
      "            surface_temperature = 5000 + 1000 * (current_mass - 0.5)\n",
      "        else:\n",
      "            surface_temperature = 6500 + 5000 * (current_mass / 10.0) ** 0.5\n",
      "    elif relative_age < 1.2:  # Giant phase\n",
      "        surface_temperature = 4000 - 1000 * (relative_age - 1.0) * 5\n",
      "    else:  # Remnant\n",
      "        if current_mass < 1.4:  # White dwarf\n",
      "            surface_temperature = 20000 * max(0.1, 1.0 - 0.1 * (relative_age - 1.3))\n",
      "        elif current_mass < 3.0:  # Neutron star\n",
      "            surface_temperature = 1000000  # Very hot surface\n",
      "        else:  # Black hole\n",
      "            surface_temperature = 0  # No surface\n",
      "\n",
      "    # Stellar stage classification\n",
      "    if relative_age < 0.01:\n",
      "        stellar_stage = \"Protostar\"\n",
      "    elif relative_age < 1.0:\n",
      "        stellar_stage = \"Main Sequence\"\n",
      "    elif relative_age < 1.05:\n",
      "        stellar_stage = \"Subgiant\"\n",
      "    elif relative_age < 1.2:\n",
      "        stellar_stage = \"Red Giant\"\n",
      "    elif relative_age < 1.3:\n",
      "        stellar_stage = \"Asymptotic Giant Branch\"\n",
      "    else:  # Remnant\n",
      "        if current_mass < 1.4:\n",
      "            stellar_stage = \"White Dwarf\"\n",
      "        elif current_mass < 3.0:\n",
      "            stellar_stage = \"Neutron Star\"\n",
      "        else:\n",
      "            stellar_stage = \"Black Hole\"\n",
      "\n",
      "    return {\n",
      "        'current_mass': current_mass,\n",
      "        'luminosity': luminosity,\n",
      "        'surface_temperature': surface_temperature,\n",
      "        'stellar_stage': stellar_stage,\n",
      "        'main_sequence_lifetime': main_sequence_lifetime  # In billions of years\n",
      "    }\n",
      "\n"
     ]
    }
   ],
   "source": [
    "function_definition = client.get_function_source(function_name=stellar_func_name)\n",
    "\n",
    "print(function_definition)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "139489d6-872f-46cd-ae76-dea54a290990",
   "metadata": {},
   "source": [
    "## Registering the Enhanced Function to Unity Catalog\n",
    "\n",
    "After developing our enhanced stellar properties function with additional astrophysical parameters, we need to register it to Unity Catalog to make it available for AI assistants and other applications. This step highlights another powerful aspect of the Unity Catalog AI workflow - the ability to maintain multiple versions or variants of related functions within the same catalog.\n",
    "\n",
    "The registration process for our enhanced function follows the same pattern we used earlier. We call `create_python_function()` with our new enhanced function, using the same catalog and schema but a different function name. Once registered, we store the fully qualified function name for later use.\n",
    "\n",
    "What's remarkable here is how Unity Catalog AI handles this complex function:\n",
    "\n",
    "1. **Comprehensive Docstring Analysis**: It parses our detailed documentation, including the expanded return value specification with spectral types, star colors, and other new fields\n",
    "2. **Parameter Validation Logic**: It preserves all our input validation code, ensuring the function remains safe to call\n",
    "3. **Complex Calculation Preservation**: It accurately captures all of our sophisticated astrophysics calculations\n",
    "4. **Local Module Handling**: It properly manages the `math` module that's imported within the function\n",
    "\n",
    "By storing the qualified function name, we can now call this function just like our original version. This approach allows us to maintain multiple related functions that serve different needs:\n",
    "\n",
    "- The original function for basic stellar evolution modeling\n",
    "- The enhanced function for more detailed astrophysical analysis\n",
    "\n",
    "This pattern is particularly valuable for scientific computing workflows, where you might need both simplified models for quick calculations and more comprehensive models for detailed analysis.\n",
    "\n",
    "The ability to register multiple versions of related functions also supports an incremental development approach, where you can evolve your functions over time without breaking existing applications that depend on earlier versions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6b3bd901-1aff-40bf-baed-fdfc13e21b76",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_stellar_properties_enhanced(\n",
    "    initial_mass: float, metallicity: float, age: float\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Calculate comprehensive stellar characteristics based on initial mass, metallicity, and age.\n",
    "\n",
    "    This function provides an enhanced model of stellar evolution and properties.\n",
    "\n",
    "    Args:\n",
    "        initial_mass (float): Initial stellar mass in solar masses (M☉).\n",
    "            Valid range: 0.1 to 150 solar masses.\n",
    "        metallicity (float): Stellar metallicity (Z), representing the fraction\n",
    "            of stellar mass composed of elements heavier than hydrogen and helium.\n",
    "            Valid range: 0.0001 to 0.04.\n",
    "        age (float): Stellar age in billions of years.\n",
    "            Valid range: 0 to 13.8 (age of the universe).\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary of stellar characteristics including:\n",
    "            - current_mass: Remaining stellar mass in solar masses\n",
    "            - luminosity: Current stellar luminosity relative to the Sun\n",
    "            - surface_temperature: Effective surface temperature in Kelvin\n",
    "            - stellar_stage: Textual description of the star's evolutionary stage\n",
    "            - spectral_type: Harvard spectral classification with luminosity class\n",
    "            - star_color: Approximate visible color of the star\n",
    "            - main_sequence_lifetime: Expected lifetime on main sequence in billions of years\n",
    "            - habitable_zone_inner_radius: Inner boundary of potential planetary habitable zone in AU\n",
    "            - habitable_zone_outer_radius: Outer boundary of potential planetary habitable zone in AU\n",
    "            - stellar_wind_velocity: Estimated stellar wind speed in km/s\n",
    "            - magnetic_field_strength: Estimated magnetic field strength at stellar surface in Gauss\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If input parameters are outside physically meaningful ranges.\n",
    "\n",
    "    Example:\n",
    "        >>> result = calculate_stellar_properties_enhanced(2.5, 0.02, 1.0)\n",
    "        >>> print(result[\"luminosity\"])\n",
    "    \"\"\"\n",
    "    # Add local import\n",
    "    import math\n",
    "\n",
    "    # Input validation\n",
    "    if not (0.1 <= initial_mass <= 150):\n",
    "        raise ValueError(\"Initial mass must be between 0.1 and 150 solar masses\")\n",
    "\n",
    "    if not (0.0001 <= metallicity <= 0.04):\n",
    "        raise ValueError(\"Metallicity must be between 0.0001 and 0.04\")\n",
    "\n",
    "    if not (0 <= age <= 13.8):\n",
    "        raise ValueError(\"Stellar age must be between 0 and 13.8 billion years\")\n",
    "\n",
    "    # Calculate stellar lifetime (in billions of years)\n",
    "    # Massive stars burn fuel much faster - lifetime scales roughly as M^-2.5\n",
    "    main_sequence_lifetime = 10.0 * (initial_mass / 1.0) ** -2.5\n",
    "\n",
    "    # Adjust lifetime based on metallicity (metal-poor stars live slightly longer)\n",
    "    metallicity_factor = (metallicity / 0.02) ** 0.3\n",
    "    main_sequence_lifetime = main_sequence_lifetime / metallicity_factor\n",
    "\n",
    "    # Calculate relative age (how far through its lifetime the star is)\n",
    "    relative_age = min(age / main_sequence_lifetime, 5.0)  # Cap at 5x lifetime for remnants\n",
    "\n",
    "    # Determine current mass based on evolutionary stage\n",
    "    if relative_age < 0.1:  # Very early evolution\n",
    "        # Minimal mass loss during early stages\n",
    "        current_mass = initial_mass * (1.0 - 0.01 * relative_age)\n",
    "    elif relative_age < 1.0:  # Main sequence\n",
    "        # Gradual mass loss during main sequence\n",
    "        # More massive stars lose more (Wolf-Rayet stars, etc)\n",
    "        main_sequence_loss = 0.05 * (initial_mass / 10.0) ** 1.5\n",
    "        current_mass = initial_mass * (1.0 - main_sequence_loss * relative_age)\n",
    "    elif relative_age < 1.1:  # End of main sequence/subgiant\n",
    "        # Increased mass loss as star expands\n",
    "        current_mass = initial_mass * (0.95 - 0.1 * (relative_age - 1.0))\n",
    "    elif relative_age < 1.3:  # Red giant/AGB phase - dramatic mass loss\n",
    "        # Up to 50-80% mass loss during giant phases\n",
    "        giant_phase_factor = 0.5 if initial_mass < 8.0 else 0.3\n",
    "        current_mass = initial_mass * (0.85 - giant_phase_factor * (relative_age - 1.1) / 0.2)\n",
    "    else:  # Remnant\n",
    "        # Determine final remnant mass based on initial mass\n",
    "        if initial_mass < 8.0:  # White dwarf territory\n",
    "            # Most stars become 0.5-1.0 M☉ white dwarfs\n",
    "            current_mass = 0.5 + 0.1 * initial_mass\n",
    "            current_mass = min(current_mass, 1.4)  # Chandrasekhar limit\n",
    "        elif initial_mass < 20.0:  # Neutron star territory\n",
    "            # Stars 8-20 M☉ tend to make ~1.4-2.0 M☉ neutron stars\n",
    "            current_mass = 1.4 + 0.03 * (initial_mass - 8.0)\n",
    "        else:  # Black hole territory\n",
    "            # Stars above 20 M☉ make black holes roughly 10-20% of initial mass\n",
    "            current_mass = max(3.0, 0.2 * initial_mass)\n",
    "\n",
    "    # Luminosity calculation\n",
    "    if relative_age < 1.0:  # Main sequence luminosity\n",
    "        if current_mass < 0.43:\n",
    "            luminosity = current_mass**2.3\n",
    "        elif current_mass < 2:\n",
    "            luminosity = current_mass**4\n",
    "        elif current_mass < 20:\n",
    "            luminosity = current_mass**3.5\n",
    "        else:\n",
    "            luminosity = current_mass**3.0\n",
    "    elif relative_age < 1.3:  # Giant phase luminosity spike\n",
    "        luminosity = 100 * current_mass**2\n",
    "    else:  # Remnant luminosity\n",
    "        if current_mass < 1.4:  # White dwarf\n",
    "            # White dwarfs cool over time\n",
    "            cooling_factor = max(0.001, 1.0 - (relative_age - 1.3))\n",
    "            luminosity = 0.1 * cooling_factor\n",
    "        elif current_mass < 3.0:  # Neutron star\n",
    "            luminosity = 0.01  # Very low luminosity\n",
    "        else:  # Black hole\n",
    "            # Accretion-dependent, but generally near-zero for isolated black holes\n",
    "            luminosity = 0.0\n",
    "\n",
    "    # Surface temperature estimation\n",
    "    if relative_age < 0.1:  # Protostar\n",
    "        surface_temperature = 3000 + 1000 * relative_age * 10\n",
    "    elif relative_age < 1.0:  # Main sequence\n",
    "        if current_mass < 0.5:\n",
    "            surface_temperature = 3000 + 1000 * current_mass\n",
    "        elif current_mass < 1.5:\n",
    "            surface_temperature = 5000 + 1000 * (current_mass - 0.5)\n",
    "        else:\n",
    "            surface_temperature = 6500 + 5000 * (current_mass / 10.0) ** 0.5\n",
    "    elif relative_age < 1.2:  # Giant phase\n",
    "        surface_temperature = 4000 - 1000 * (relative_age - 1.0) * 5\n",
    "    else:  # Remnant\n",
    "        if current_mass < 1.4:  # White dwarf\n",
    "            surface_temperature = 20000 * max(0.1, 1.0 - 0.1 * (relative_age - 1.3))\n",
    "        elif current_mass < 3.0:  # Neutron star\n",
    "            surface_temperature = 1000000  # Very hot surface\n",
    "        else:  # Black hole\n",
    "            surface_temperature = 0  # No surface\n",
    "\n",
    "    # Stellar stage classification\n",
    "    if relative_age < 0.01:\n",
    "        stellar_stage = \"Protostar\"\n",
    "    elif relative_age < 1.0:\n",
    "        stellar_stage = \"Main Sequence\"\n",
    "    elif relative_age < 1.05:\n",
    "        stellar_stage = \"Subgiant\"\n",
    "    elif relative_age < 1.2:\n",
    "        stellar_stage = \"Red Giant\"\n",
    "    elif relative_age < 1.3:\n",
    "        stellar_stage = \"Asymptotic Giant Branch\"\n",
    "    else:  # Remnant\n",
    "        if current_mass < 1.4:\n",
    "            stellar_stage = \"White Dwarf\"\n",
    "        elif current_mass < 3.0:\n",
    "            stellar_stage = \"Neutron Star\"\n",
    "        else:\n",
    "            stellar_stage = \"Black Hole\"\n",
    "\n",
    "    # Enhanced features\n",
    "\n",
    "    # Habitable zone calculation (only meaningful for main sequence stars)\n",
    "    # Based on stellar luminosity using simplified model from Kopparapu et al. 2013\n",
    "    if stellar_stage == \"Main Sequence\" and surface_temperature > 2600:\n",
    "        # Convert luminosity to solar units for habitable zone calculation\n",
    "        # Inner edge (runaway greenhouse limit)\n",
    "        habitable_zone_inner_radius = 0.75 * (luminosity**0.5)\n",
    "        # Outer edge (maximum greenhouse limit)\n",
    "        habitable_zone_outer_radius = 1.77 * (luminosity**0.5)\n",
    "    else:\n",
    "        habitable_zone_inner_radius = None\n",
    "        habitable_zone_outer_radius = None\n",
    "\n",
    "    # Stellar wind velocity estimation in km/s\n",
    "    # More sophisticated model based on stellar type and surface temperature\n",
    "    if stellar_stage in [\"Main Sequence\", \"Subgiant\"]:\n",
    "        if surface_temperature > 30000:  # O stars\n",
    "            stellar_wind_velocity = 2000 + 500 * (surface_temperature / 40000)\n",
    "        elif surface_temperature > 10000:  # B and A stars\n",
    "            stellar_wind_velocity = 1000 + 1000 * (surface_temperature / 30000)\n",
    "        elif surface_temperature > 5800:  # F and G stars\n",
    "            stellar_wind_velocity = 450 + 550 * ((surface_temperature - 5800) / 4200)\n",
    "        else:  # K and M stars\n",
    "            stellar_wind_velocity = 300 + 150 * ((surface_temperature - 3000) / 2800)\n",
    "    elif stellar_stage in [\"Red Giant\", \"Asymptotic Giant Branch\"]:\n",
    "        stellar_wind_velocity = 30 + 10 * (luminosity**0.25)  # Slow dense winds\n",
    "    elif stellar_stage == \"White Dwarf\":\n",
    "        stellar_wind_velocity = 4000  # Hot fast winds\n",
    "    elif stellar_stage == \"Neutron Star\":\n",
    "        stellar_wind_velocity = 10000  # Extremely fast winds if active\n",
    "    else:\n",
    "        stellar_wind_velocity = 0\n",
    "\n",
    "    # Magnetic field strength estimation in Gauss\n",
    "    # Based on empirical relationships for different stellar types\n",
    "    if stellar_stage == \"Main Sequence\":\n",
    "        if surface_temperature > 15000:  # Hot massive stars\n",
    "            # Hot stars have weaker fields due to fully radiative envelopes\n",
    "            magnetic_field_strength = 100 * (1 - 0.9 * (age / main_sequence_lifetime))\n",
    "        elif surface_temperature > 5500:  # Sun-like stars\n",
    "            # Solar-type stars have moderate fields that decay with age\n",
    "            magnetic_field_strength = 500 * (1 - 0.7 * (age / main_sequence_lifetime))\n",
    "        else:  # M dwarfs and cool stars\n",
    "            # Cool stars can have very strong fields\n",
    "            rotation_factor = math.exp(-0.5 * age / main_sequence_lifetime)\n",
    "            magnetic_field_strength = 2000 * rotation_factor\n",
    "    elif stellar_stage in [\"Red Giant\", \"Asymptotic Giant Branch\"]:\n",
    "        magnetic_field_strength = 10  # Generally weak fields\n",
    "    elif stellar_stage == \"White Dwarf\":\n",
    "        # Some white dwarfs have extremely strong fields\n",
    "        if current_mass > 1.0:\n",
    "            magnetic_field_strength = 10**6  # Magnetic white dwarf\n",
    "        else:\n",
    "            magnetic_field_strength = 10**3  # Typical white dwarf\n",
    "    elif stellar_stage == \"Neutron Star\":\n",
    "        magnetic_field_strength = 10**12  # Typical neutron star field\n",
    "    else:\n",
    "        magnetic_field_strength = 0\n",
    "\n",
    "    # Chemical enrichment potential (0-100 scale)\n",
    "    # Indicates how much a star enriches its environment with metals\n",
    "    if stellar_stage == \"Main Sequence\":\n",
    "        if initial_mass > 8:\n",
    "            chemical_enrichment = 80 + 20 * (\n",
    "                initial_mass / 40\n",
    "            )  # Massive stars produce many elements\n",
    "        else:\n",
    "            chemical_enrichment = 20 + 10 * initial_mass  # Lower mass stars produce fewer elements\n",
    "    elif stellar_stage in [\"Red Giant\", \"Asymptotic Giant Branch\"]:\n",
    "        chemical_enrichment = 70 + 30 * min(\n",
    "            initial_mass / 8, 1.0\n",
    "        )  # AGB stars produce s-process elements\n",
    "    elif stellar_stage in [\"Neutron Star\", \"Black Hole\"]:\n",
    "        # Past supernova has enriched environment\n",
    "        chemical_enrichment = 95 + 5 * min(initial_mass / 30, 1.0)\n",
    "    else:\n",
    "        chemical_enrichment = 10  # White dwarfs and other stages contribute little\n",
    "\n",
    "    # Determine spectral type and color for main sequence stars\n",
    "    spectral_type = None\n",
    "    star_color = None\n",
    "\n",
    "    if stellar_stage == \"Main Sequence\":\n",
    "        # Spectral classification based on temperature\n",
    "        if surface_temperature >= 30000:\n",
    "            spectral_type = \"O\"\n",
    "            star_color = \"Blue\"\n",
    "        elif surface_temperature >= 10000:\n",
    "            spectral_type = \"B\"\n",
    "            star_color = \"Blue-white\"\n",
    "        elif surface_temperature >= 7500:\n",
    "            spectral_type = \"A\"\n",
    "            star_color = \"White\"\n",
    "        elif surface_temperature >= 6000:\n",
    "            spectral_type = \"F\"\n",
    "            star_color = \"Yellow-white\"\n",
    "        elif surface_temperature >= 5200:\n",
    "            spectral_type = \"G\"\n",
    "            star_color = \"Yellow\"\n",
    "        elif surface_temperature >= 3700:\n",
    "            spectral_type = \"K\"\n",
    "            star_color = \"Orange\"\n",
    "        else:\n",
    "            spectral_type = \"M\"\n",
    "            star_color = \"Red\"\n",
    "\n",
    "        # Add luminosity class (V for main sequence)\n",
    "        spectral_type += \"V\"\n",
    "\n",
    "    elif stellar_stage == \"Subgiant\":\n",
    "        # Subgiants are luminosity class IV\n",
    "        if surface_temperature >= 10000:\n",
    "            spectral_type = \"B IV\"\n",
    "            star_color = \"Blue-white\"\n",
    "        elif surface_temperature >= 7500:\n",
    "            spectral_type = \"A IV\"\n",
    "            star_color = \"White\"\n",
    "        elif surface_temperature >= 6000:\n",
    "            spectral_type = \"F IV\"\n",
    "            star_color = \"Yellow-white\"\n",
    "        elif surface_temperature >= 5200:\n",
    "            spectral_type = \"G IV\"\n",
    "            star_color = \"Yellow\"\n",
    "        elif surface_temperature >= 3700:\n",
    "            spectral_type = \"K IV\"\n",
    "            star_color = \"Orange\"\n",
    "        else:\n",
    "            spectral_type = \"M IV\"\n",
    "            star_color = \"Red\"\n",
    "\n",
    "    elif stellar_stage in [\"Red Giant\", \"Asymptotic Giant Branch\"]:\n",
    "        # Giants are luminosity class III\n",
    "        if surface_temperature >= 5200:\n",
    "            spectral_type = \"G III\"\n",
    "            star_color = \"Yellow\"\n",
    "        elif surface_temperature >= 3700:\n",
    "            spectral_type = \"K III\"\n",
    "            star_color = \"Orange\"\n",
    "        else:\n",
    "            spectral_type = \"M III\"\n",
    "            star_color = \"Red\"\n",
    "\n",
    "    elif stellar_stage == \"White Dwarf\":\n",
    "        spectral_type = \"D\"\n",
    "        star_color = \"White\"\n",
    "\n",
    "    elif stellar_stage == \"Neutron Star\":\n",
    "        spectral_type = \"Neutron Star\"\n",
    "        star_color = None\n",
    "\n",
    "    elif stellar_stage == \"Black Hole\":\n",
    "        spectral_type = \"Black Hole\"\n",
    "        star_color = None\n",
    "\n",
    "    # Create return dictionary with all properties\n",
    "    result = {\n",
    "        \"current_mass\": current_mass,\n",
    "        \"luminosity\": luminosity,\n",
    "        \"surface_temperature\": surface_temperature,\n",
    "        \"stellar_stage\": stellar_stage,\n",
    "        \"spectral_type\": spectral_type,\n",
    "        \"star_color\": star_color,\n",
    "        \"main_sequence_lifetime\": main_sequence_lifetime,\n",
    "        \"habitable_zone_inner_radius\": habitable_zone_inner_radius,\n",
    "        \"habitable_zone_outer_radius\": habitable_zone_outer_radius,\n",
    "        \"stellar_wind_velocity\": stellar_wind_velocity,\n",
    "        \"magnetic_field_strength\": magnetic_field_strength,\n",
    "        \"chemical_enrichment_potential\": chemical_enrichment,\n",
    "    }\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "74f22e9b-fadc-44bd-a1c4-6dd9f2cb60b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.create_python_function(\n",
    "    func=calculate_stellar_properties_enhanced,\n",
    "    catalog=CATALOG,\n",
    "    schema=SCHEMA,\n",
    "    replace=True,\n",
    ")\n",
    "\n",
    "stellar_func_name_enhanced = f\"{CATALOG}.{SCHEMA}.calculate_stellar_properties_enhanced\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cbbfaf9-0224-4ae4-9846-3416e5a493c1",
   "metadata": {},
   "source": [
    "## Executing the Enhanced Function with Test Scenarios\n",
    "\n",
    "After registering our enhanced stellar function with Unity Catalog, we need to validate it with a comprehensive set of test scenarios. This testing phase is critical to ensure our function behaves as expected across the parameter space that AI assistants might explore.\n",
    "\n",
    "For our enhanced stellar properties function, we've selected four representative scenarios that showcase different stellar types and evolutionary stages:\n",
    "\n",
    "1. **Sun-like Star**: A G-type main sequence star with solar metallicity and middle age (1.0 M☉)\n",
    "2. **Massive O-Type Star**: A young, high-mass star that will eventually form a black hole (25.0 M☉)\n",
    "3. **Low-mass Red Dwarf**: A common, long-lived, cool star (0.3 M☉)\n",
    "4. **Evolved Massive Star**: An old, massive star that has already formed a black hole (90.5 M☉)\n",
    "\n",
    "These scenarios test our function across orders of magnitude in stellar mass and through different evolutionary stages, from young main sequence stars to evolved stellar remnants.\n",
    "\n",
    "For each test case, we:\n",
    "1. Call the function via `execute_function()` with the appropriate parameters\n",
    "2. Parse the returned string representation into a Python dictionary\n",
    "3. Display all the calculated stellar properties for analysis\n",
    "\n",
    "This systematic testing approach helps catch edge cases and ensures that both fundamental properties (mass, luminosity, temperature) and our enhanced properties (spectral type, habitable zone, magnetic field strength) are calculated correctly for all stellar types.\n",
    "\n",
    "Testing different evolutionary stages is particularly important for functions that model time-dependent processes. Our test set includes stars in various life stages, from young stars to those that have already completed their evolution, ensuring our function handles the entire stellar lifecycle appropriately.\n",
    "\n",
    "This validation step provides confidence that when AI assistants or data scientists use our function, it will produce scientifically sound results across the entire parameter space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f58fbd9b-e430-44ae-8946-d9281d9bfd6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Stellar Model for 1.0 M☉ star:\n",
      "\tcurrent_mass: 0.9992726761381613\n",
      "\tluminosity: 0.9970938770139077\n",
      "\tsurface_temperature: 5499.272676138161\n",
      "\tstellar_stage: Main Sequence\n",
      "\tspectral_type: GV\n",
      "\tstar_color: Yellow\n",
      "\tmain_sequence_lifetime: 10.0\n",
      "\thabitable_zone_inner_radius: 0.748909410957242\n",
      "\thabitable_zone_outer_radius: 1.767426209859091\n",
      "\tstellar_wind_velocity: 433.88960765025865\n",
      "\tmagnetic_field_strength: 1589.067205006668\n",
      "\tchemical_enrichment_potential: 30.0\n",
      "\n",
      "Stellar Model for 25.0 M☉ star:\n",
      "\tcurrent_mass: 22.16718363676904\n",
      "\tluminosity: 10892.600037747503\n",
      "\tsurface_temperature: 13944.323951301596\n",
      "\tstellar_stage: Main Sequence\n",
      "\tspectral_type: BV\n",
      "\tstar_color: Blue-white\n",
      "\tmain_sequence_lifetime: 0.0034884427442219237\n",
      "\thabitable_zone_inner_radius: 78.27571476028163\n",
      "\thabitable_zone_outer_radius: 184.73068683426465\n",
      "\tstellar_wind_velocity: 1464.81079837672\n",
      "\tmagnetic_field_strength: 299.3373974219747\n",
      "\tchemical_enrichment_potential: 92.5\n",
      "\n",
      "Stellar Model for 0.3 M☉ star:\n",
      "\tcurrent_mass: 0.2999397973798912\n",
      "\tluminosity: 0.06268713412999312\n",
      "\tsurface_temperature: 3200.675400362578\n",
      "\tstellar_stage: Main Sequence\n",
      "\tspectral_type: MV\n",
      "\tstar_color: Red\n",
      "\tmain_sequence_lifetime: 498.3171819730835\n",
      "\thabitable_zone_inner_radius: 0.18778049139386427\n",
      "\thabitable_zone_outer_radius: 0.44316195968951966\n",
      "\tstellar_wind_velocity: 310.7504678765667\n",
      "\tmagnetic_field_strength: 1980.0328006257218\n",
      "\tchemical_enrichment_potential: 23.0\n",
      "\n",
      "Stellar Model for 90.5 M☉ star:\n",
      "\tcurrent_mass: 18.1\n",
      "\tluminosity: 0.0\n",
      "\tsurface_temperature: 0\n",
      "\tstellar_stage: Black Hole\n",
      "\tspectral_type: Black Hole\n",
      "\tstar_color: None\n",
      "\tmain_sequence_lifetime: 0.00019453446251488867\n",
      "\thabitable_zone_inner_radius: None\n",
      "\thabitable_zone_outer_radius: None\n",
      "\tstellar_wind_velocity: 0\n",
      "\tmagnetic_field_strength: 0\n",
      "\tchemical_enrichment_potential: 100.0\n"
     ]
    }
   ],
   "source": [
    "scenarios_enhanced = [\n",
    "    # Sun-like star\n",
    "    (1.0, 0.02, 4.6),\n",
    "    # Massive star\n",
    "    (25.0, 0.015, 0.002),\n",
    "    # Low-mass star\n",
    "    (0.3, 0.001, 10.0),\n",
    "    # Old Massive Star that has evolved to a singularity\n",
    "    (90.5, 0.005, 5.0),\n",
    "]\n",
    "\n",
    "for mass, metallicity, age in scenarios_enhanced:\n",
    "    result = client.execute_function(\n",
    "        function_name=stellar_func_name_enhanced,\n",
    "        parameters={\n",
    "            \"initial_mass\": mass,\n",
    "            \"metallicity\": metallicity,\n",
    "            \"age\": age,\n",
    "        },\n",
    "    )\n",
    "    print(f\"\\nStellar Model for {mass} M☉ star:\")\n",
    "    for key, value in ast.literal_eval(result.value).items():\n",
    "        print(f\"\\t{key}: {value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a27fc266-6036-484e-84c8-cd73e38177a0",
   "metadata": {},
   "source": [
    "## Testing a function locally\n",
    "\n",
    "There can be times when all you want is to retrieve a function definition from Unity Catalog and have it available in your current process, directly as a callable. \n",
    "\n",
    "The `get_function_as_callable` API allows the flexibility for several options:\n",
    "\n",
    "- **Default behavior (non-Jupyter)**: Registers the original function name into the global namespace. You can call the function as if you had defined its implementation in the REPL directly.\n",
    "- **Global namespace in Jupyter)**: You can directly define the return value of the call to `get_function_as_callable` to the active globals namespace **or** just call the return value as a callable directly (shown in the next cell below)\n",
    "\n",
    "To set the callable to the global namespace in Jupyter notebooks:\n",
    "\n",
    "```python\n",
    "retrieved_func = client.get_function_as_callable(<catalog.schema.func_name>)\n",
    "globals()[\"calculate_stellar_properties\"] = retrieved_func\n",
    "```\n",
    "\n",
    "- **Custom namespace**: If you have a namespace that you would like the callable to be registered in, you can submit a `SimpleNamespace` instance (from the `types` module) to restrict namespace scoping of the loaded function.\n",
    "- **No registration**: You can set the argument `register_function` to `False`) to only allow for the return value to be used explicitly.\n",
    "\n",
    "```python\n",
    "retrieved_func = client.get_function_as_callable(<catalog.schema.func_name>, register_function=False)\n",
    "# The only way to use this function is by directly calling the return instance.\n",
    "retrieved_func(initial_mass=1.3, metallicity=0.0002, age=0.65)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "63ddf723-1e51-41a1-af02-c521c1077abe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'current_mass': 19.840000000000003,\n",
       " 'luminosity': 0.0,\n",
       " 'surface_temperature': 0,\n",
       " 'stellar_stage': 'Black Hole',\n",
       " 'main_sequence_lifetime': 0.0003085594253277336}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retrieved_func = client.get_function_as_callable(stellar_func_name)\n",
    "\n",
    "retrieved_func(initial_mass=99.2, metallicity=0.0005, age=0.005)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44a6781-ea1b-40ac-aa74-ec32ff39356e",
   "metadata": {},
   "source": [
    "## Creating Independent Black Hole Physics Functions\n",
    "\n",
    "To demonstrate the power of function wrapping, let's shift our focus from stellar astrophysics to black hole physics. In this section, we'll create a set of independent functions that calculate various properties of black holes.\n",
    "\n",
    "Each of these functions represents a distinct physical calculation related to black holes:\n",
    "\n",
    "1. **Schwarzschild Radius**: Calculates the event horizon radius (the \"point of no return\") for a black hole of a given mass\n",
    "2. **Hawking Temperature**: Determines the quantum mechanical thermal radiation emitted by black holes\n",
    "3. **Event Horizon Area**: Computes the surface area of the event horizon, which relates to black hole entropy\n",
    "4. **Accretion Luminosity**: Estimates the energy output from matter falling into a black hole\n",
    "5. **Eddington Luminosity**: Calculates the theoretical maximum luminosity a black hole can achieve\n",
    "\n",
    "These functions are designed to be independently useful while also serving as building blocks for more comprehensive analyses. Each function:\n",
    "\n",
    "- Accepts clear, well-defined parameters with validation\n",
    "- Performs a specific calculation based on established physics equations\n",
    "- Returns a single, specific physical property\n",
    "- Includes comprehensive documentation explaining its purpose\n",
    "\n",
    "This modular approach to function development follows scientific best practices by isolating specific physical phenomena into discrete, testable units. Each function can be validated against known solutions and used independently when only a specific property is needed.\n",
    "\n",
    "When registering these functions with Unity Catalog, they become individually accessible to AI assistants and other applications. This granularity allows AI systems to select exactly the function they need based on user queries - a question about black hole temperature would invoke just the Hawking temperature function, while a question about accretion would use the relevant accretion function.\n",
    "\n",
    "This approach also makes the functions more maintainable, as updates to one physical calculation don't affect the others. For example, if a more accurate formula for Hawking radiation is developed, we can update just that function without touching the others."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3b85b2b9-4238-4eef-b4c8-357629d0c3a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def schwarzschild_radius(mass_solar_masses: float) -> float:\n",
    "    \"\"\"\n",
    "    Calculate the Schwarzschild radius (event horizon) of a black hole.\n",
    "\n",
    "    Args:\n",
    "        mass_solar_masses (float): Mass of the black hole in solar masses.\n",
    "            Valid range: Greater than 0.\n",
    "\n",
    "    Returns:\n",
    "        float: Schwarzschild radius in kilometers.\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If mass is not positive.\n",
    "\n",
    "    Example:\n",
    "        >>> schwarzschild_radius(10)\n",
    "        29.5327\n",
    "    \"\"\"\n",
    "    if mass_solar_masses <= 0:\n",
    "        raise ValueError(\"Mass must be positive\")\n",
    "\n",
    "    # Convert solar masses to kg\n",
    "    mass_kg = mass_solar_masses * 1.989e30\n",
    "\n",
    "    # Constants\n",
    "    G = 6.67430e-11  # Gravitational constant in m^3 kg^-1 s^-2\n",
    "    c = 299792458  # Speed of light in m/s\n",
    "\n",
    "    # Calculate Schwarzschild radius in meters\n",
    "    r_s_meters = 2 * G * mass_kg / c**2\n",
    "\n",
    "    # Convert to kilometers\n",
    "    r_s_km = r_s_meters / 1000\n",
    "\n",
    "    return r_s_km\n",
    "\n",
    "\n",
    "def hawking_temperature(mass_solar_masses: float) -> float:\n",
    "    \"\"\"\n",
    "    Calculate the Hawking radiation temperature of a black hole.\n",
    "\n",
    "    Args:\n",
    "        mass_solar_masses (float): Mass of the black hole in solar masses.\n",
    "            Valid range: Greater than 0.\n",
    "\n",
    "    Returns:\n",
    "        float: Hawking radiation temperature in nanokelvin.\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If mass is not positive.\n",
    "\n",
    "    Example:\n",
    "        >>> hawking_temperature(10)\n",
    "        6.1707e-10\n",
    "    \"\"\"\n",
    "    import math\n",
    "\n",
    "    if mass_solar_masses <= 0:\n",
    "        raise ValueError(\"Mass must be positive\")\n",
    "\n",
    "    # Convert solar masses to kg\n",
    "    mass_kg = mass_solar_masses * 1.989e30\n",
    "\n",
    "    # Constants\n",
    "    h_bar = 1.054571817e-34  # Reduced Planck constant in J·s\n",
    "    c = 299792458  # Speed of light in m/s\n",
    "    G = 6.67430e-11  # Gravitational constant in m^3 kg^-1 s^-2\n",
    "    k_B = 1.380649e-23  # Boltzmann constant in J/K\n",
    "\n",
    "    # Calculate Hawking temperature in Kelvin\n",
    "    T = (h_bar * c**3) / (8 * math.pi * G * k_B * mass_kg)\n",
    "\n",
    "    # Convert to nanokelvin for more readable numbers\n",
    "    T_nK = T * 1e9\n",
    "\n",
    "    return T_nK\n",
    "\n",
    "\n",
    "def event_horizon_area(mass_solar_masses: float) -> float:\n",
    "    \"\"\"\n",
    "    Calculate the surface area of a black hole's event horizon.\n",
    "\n",
    "    Args:\n",
    "        mass_solar_masses (float): Mass of the black hole in solar masses.\n",
    "            Valid range: Greater than 0.\n",
    "\n",
    "    Returns:\n",
    "        float: Event horizon surface area in square kilometers.\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If mass is not positive.\n",
    "\n",
    "    Example:\n",
    "        >>> event_horizon_area(10)\n",
    "        10955.5\n",
    "    \"\"\"\n",
    "    import math\n",
    "\n",
    "    if mass_solar_masses <= 0:\n",
    "        raise ValueError(\"Mass must be positive\")\n",
    "\n",
    "    # Constants\n",
    "    G = 6.67430e-11  # Gravitational constant in m^3 kg^-1 s^-2\n",
    "    c = 299792458  # Speed of light in m/s\n",
    "\n",
    "    # Convert solar masses to kg\n",
    "    mass_kg = mass_solar_masses * 1.989e30\n",
    "\n",
    "    # Calculate Schwarzschild radius in meters\n",
    "    r_s_meters = 2 * G * mass_kg / c**2\n",
    "\n",
    "    # Convert to kilometers\n",
    "    r_s_km = r_s_meters / 1000\n",
    "\n",
    "    # Calculate area in square kilometers\n",
    "    area = 4 * math.pi * r_s_km**2\n",
    "\n",
    "    return area\n",
    "\n",
    "\n",
    "def accretion_luminosity(\n",
    "    mass_solar_masses: float, accretion_rate_solar_masses_per_year: float, efficiency: float = 0.1\n",
    ") -> float:\n",
    "    \"\"\"\n",
    "    Calculate the luminosity from matter accreting onto a black hole.\n",
    "\n",
    "    Args:\n",
    "        mass_solar_masses (float): Mass of the black hole in solar masses.\n",
    "            Valid range: Greater than 0.\n",
    "        accretion_rate_solar_masses_per_year (float): Rate at which matter falls\n",
    "            into the black hole in solar masses per year.\n",
    "            Valid range: Greater than or equal to 0.\n",
    "        efficiency (float, optional): Efficiency of converting rest mass to radiation.\n",
    "            Default is 0.1 (10%). Valid range: 0 to 1.\n",
    "\n",
    "    Returns:\n",
    "        float: Luminosity in units of solar luminosity.\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If parameters are outside valid ranges.\n",
    "\n",
    "    Example:\n",
    "        >>> accretion_luminosity(10, 1e-9)\n",
    "        5.7e4\n",
    "    \"\"\"\n",
    "    if mass_solar_masses <= 0:\n",
    "        raise ValueError(\"Mass must be positive\")\n",
    "    if accretion_rate_solar_masses_per_year < 0:\n",
    "        raise ValueError(\"Accretion rate cannot be negative\")\n",
    "    if not 0 <= efficiency <= 1:\n",
    "        raise ValueError(\"Efficiency must be between 0 and 1\")\n",
    "\n",
    "    # Constants\n",
    "    c = 299792458  # Speed of light in m/s\n",
    "    L_sun = 3.828e26  # Solar luminosity in watts\n",
    "    M_sun = 1.989e30  # Solar mass in kg\n",
    "\n",
    "    # Convert solar masses per year to kg/s\n",
    "    accretion_rate_kg_s = accretion_rate_solar_masses_per_year * M_sun / (365.25 * 24 * 3600)\n",
    "\n",
    "    # Calculate luminosity in watts\n",
    "    L_watts = efficiency * accretion_rate_kg_s * c**2\n",
    "\n",
    "    # Convert to solar luminosity\n",
    "    L_solar = L_watts / L_sun\n",
    "\n",
    "    return L_solar\n",
    "\n",
    "\n",
    "def eddington_luminosity(mass_solar_masses: float) -> float:\n",
    "    \"\"\"\n",
    "    Calculate the Eddington luminosity of a black hole.\n",
    "\n",
    "    The Eddington luminosity is the maximum luminosity a black hole can achieve\n",
    "    through accretion while maintaining hydrostatic equilibrium, where radiation\n",
    "    pressure balances gravitational attraction.\n",
    "\n",
    "    Args:\n",
    "        mass_solar_masses (float): Mass of the black hole in solar masses.\n",
    "            Valid range: Greater than 0.\n",
    "\n",
    "    Returns:\n",
    "        float: Eddington luminosity in units of solar luminosity.\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If mass is not positive.\n",
    "\n",
    "    Example:\n",
    "        >>> eddington_luminosity(10)\n",
    "        3.2e38\n",
    "    \"\"\"\n",
    "    import math\n",
    "\n",
    "    if mass_solar_masses <= 0:\n",
    "        raise ValueError(\"Mass must be positive\")\n",
    "\n",
    "    # Constants\n",
    "    G = 6.67430e-11  # Gravitational constant in m^3 kg^-1 s^-2\n",
    "    c = 299792458  # Speed of light in m/s\n",
    "    thomson_cross_section = 6.65e-29  # Thomson cross-section for electron in m²\n",
    "    m_p = 1.67262192e-27  # Proton mass in kg\n",
    "    L_sun = 3.828e26  # Solar luminosity in watts\n",
    "\n",
    "    # Convert solar masses to kg\n",
    "    mass_kg = mass_solar_masses * 1.989e30\n",
    "\n",
    "    # Calculate Eddington luminosity in watts\n",
    "    L_edd_watts = 4 * math.pi * G * mass_kg * m_p * c / thomson_cross_section\n",
    "\n",
    "    # Convert to solar luminosity\n",
    "    L_edd_solar = L_edd_watts / L_sun\n",
    "\n",
    "    return L_edd_solar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2099330c-8056-4a90-944a-65627779836e",
   "metadata": {},
   "source": [
    "## Testing Individual Functions Before Wrapping\n",
    "\n",
    "Before consolidating our functions into a wrapper, it's crucial to register and test each component function independently. This step ensures that each calculation works correctly on its own and helps isolate any issues before they propagate to more complex functions.\n",
    "\n",
    "Unity Catalog AI makes this process straightforward with the `create_python_function` API, which allows us to register each function with just a single call:\n",
    "\n",
    "```python\n",
    "for func in [schwarzschild_radius, hawking_temperature, event_horizon_area, \n",
    "             accretion_luminosity, eddington_luminosity]:\n",
    "    client.create_python_function(catalog=CATALOG, schema=SCHEMA, func=func, replace=True)\n",
    "```\n",
    "\n",
    "This approach has several key benefits for the development workflow:\n",
    "\n",
    "1. **Incremental Testing**: You can validate each specialized calculation separately\n",
    "2. **Independent Accessibility**: AI agents and other systems can call individual functions when they only need specific calculations\n",
    "3. **Clearer Debugging**: Issues can be pinpointed to specific functions rather than troubleshooting the entire wrapper\n",
    "4. **Versioning Flexibility**: You can update individual functions without necessarily having to update the wrapper\n",
    "\n",
    "For scientific applications like our black hole physics library, this modular approach mirrors how scientific understanding develops - specific phenomena are studied individually before being integrated into comprehensive models.\n",
    "\n",
    "Testing these functions individually also provides confidence that the wrapper function will work correctly, as it's built on validated components. This incremental validation approach is a best practice for complex scientific and analytical workflows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "daa88fb2-f74c-46bd-948d-c12ef63d438c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for func in [\n",
    "    schwarzschild_radius,\n",
    "    hawking_temperature,\n",
    "    event_horizon_area,\n",
    "    accretion_luminosity,\n",
    "    eddington_luminosity,\n",
    "]:\n",
    "    client.create_python_function(catalog=CATALOG, schema=SCHEMA, func=func, replace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f9be97c1-fca8-425e-b9d0-8641bb3d9b7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value=5.8888250014088315e+47, truncated=None)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.execute_function(\n",
    "    function_name=f\"{CATALOG}.{SCHEMA}.schwarzschild_radius\",\n",
    "    parameters={\"mass_solar_masses\": 1.9934234e47},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "dd377f86-16e6-4303-8500-5c2f26e23a97",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value=3.0943901394108388e-46, truncated=None)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.execute_function(\n",
    "    function_name=f\"{CATALOG}.{SCHEMA}.hawking_temperature\",\n",
    "    parameters={\"mass_solar_masses\": 1.9934234e47},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e443f463-61f1-47d8-bbb0-4eaead50871f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value=4.009071789077358e+296, truncated=None)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.execute_function(\n",
    "    function_name=f\"{CATALOG}.{SCHEMA}.event_horizon_area\",\n",
    "    parameters={\"mass_solar_masses\": 1.912e147},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "850554c2-b268-4a88-b848-c0d14f8324b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value=13022158135851.205, truncated=None)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.execute_function(\n",
    "    function_name=f\"{CATALOG}.{SCHEMA}.accretion_luminosity\",\n",
    "    parameters={\n",
    "        \"mass_solar_masses\": 1.912e147,\n",
    "        \"accretion_rate_solar_masses_per_year\": 1.1,\n",
    "        \"efficiency\": 0.8,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9accbcb1-7763-4704-b76c-9d5eeecf2268",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionExecutionResult(error=None, format='SCALAR', value=164237.2010417032, truncated=None)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.execute_function(\n",
    "    function_name=f\"{CATALOG}.{SCHEMA}.eddington_luminosity\",\n",
    "    parameters={\"mass_solar_masses\": 4.998},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aede2fb-b83d-4c79-ab2d-f44e3c1798ca",
   "metadata": {},
   "source": [
    "## Function Wrapping: Creating Integrated Solutions with Minimal Code Duplication\n",
    "\n",
    "One of the most powerful features introduced in Unity Catalog AI 0.3.0 is the `create_wrapped_function` API. This capability solves a fundamental challenge in modular programming: how to combine specialized functions into more powerful composite functions without duplicating code. Let's explore how this works in the context of our black hole physics library.\n",
    "\n",
    "### The Problem: Function Dependencies in Unity Catalog\n",
    "\n",
    "Unity Catalog functions are registered independently and can't directly call each other within the catalog environment. This creates a dilemma when you need functions to build on each other's results:\n",
    "\n",
    "1. You could **duplicate logic** across functions, but this leads to maintenance nightmares\n",
    "2. You could make **client-side calls** between functions, but this adds latency and complexity\n",
    "3. You could maintain **separate codebases** for individual functions and composites, but this creates version drift\n",
    "\n",
    "### The Solution: Automatic Function Wrapping\n",
    "\n",
    "The `create_wrapped_function` API elegantly solves this problem by:\n",
    "\n",
    "1. **Automatically injecting** the source code of dependent functions\n",
    "2. **Preserving direct function calls** in your wrapper logic\n",
    "3. **Creating a single consolidated function** in Unity Catalog that contains all necessary code\n",
    "\n",
    "With this approach, you can:\n",
    "- Develop and test each component function independently\n",
    "- Combine them into powerful composite solutions\n",
    "- Maintain all functions from a single codebase\n",
    "- Avoid code duplication and synchronization challenges\n",
    "\n",
    "In our black hole physics example, the `black_hole_analyzer` wrapper directly calls the five specialized functions (`schwarzschild_radius`, `event_horizon_area`, etc.) as if they were locally available. When registered with `create_wrapped_function`, the system automatically injects all necessary function code to make a self-contained composite function.\n",
    "\n",
    "This pattern is particularly valuable for scientific and analytical workflows, where complex analyses often build upon layers of specialized calculations. The wrapper can focus on integration logic and derived calculations while leveraging the specialized expertise in each component function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "49df8364-8ac0-4239-82a1-80dfd9700bd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the wrapper function\n",
    "def black_hole_analyzer(\n",
    "    mass_solar_masses: float,\n",
    "    accretion_rate_solar_masses_per_year: float = 0.0,\n",
    "    efficiency: float = 0.1,\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Comprehensive analysis of black hole properties based on mass and accretion.\n",
    "\n",
    "    This function acts as a wrapper that calls multiple black hole physics\n",
    "    functions and combines their results into a single comprehensive report.\n",
    "    It also performs additional calculations based on the results of these functions.\n",
    "\n",
    "    Args:\n",
    "        mass_solar_masses: Mass of the black hole in solar masses.\n",
    "            Valid range: Greater than 0.\n",
    "        accretion_rate_solar_masses_per_year: Rate at which matter falls\n",
    "            into the black hole in solar masses per year. Default is 0.0 (no accretion).\n",
    "        efficiency: Efficiency of converting rest mass to radiation\n",
    "            during accretion. Default is 0.1 (10%). Valid range: 0 to 1.\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary containing multiple black hole properties:\n",
    "            - event_horizon_radius_km: Schwarzschild radius in kilometers\n",
    "            - event_horizon_area_km2: Surface area of the event horizon in square kilometers\n",
    "            - hawking_temperature_nK: Hawking radiation temperature in nanokelvin\n",
    "            - accretion_luminosity_solar: Luminosity from accretion in units of solar luminosity\n",
    "            - eddington_luminosity_solar: Maximum sustainable luminosity in solar units\n",
    "            - eddington_ratio: Ratio of accretion luminosity to Eddington luminosity\n",
    "            - time_to_evaporation_years: Estimated time until evaporation due to Hawking radiation\n",
    "            - photon_sphere_radius_km: Radius of the photon sphere where light orbits the black hole\n",
    "            - innermost_stable_orbit_km: Radius of the innermost stable circular orbit\n",
    "\n",
    "    Example:\n",
    "        >>> result = black_hole_analyzer(10, 1e-9)\n",
    "        >>> print(f\"Event horizon: {result['event_horizon_radius_km']:.2f} km\")\n",
    "        Event horizon: 29.53 km\n",
    "    \"\"\"\n",
    "    # The function bodies will be automatically injected by create_wrapped_function\n",
    "    # so we can directly call the functions that will be wrapped\n",
    "\n",
    "    # Get results from the independent functions\n",
    "    event_horizon_radius_km = schwarzschild_radius(mass_solar_masses)\n",
    "    event_horizon_area_km2 = event_horizon_area(mass_solar_masses)\n",
    "    hawking_temperature_nK = hawking_temperature(mass_solar_masses)\n",
    "    accretion_luminosity_solar = accretion_luminosity(\n",
    "        mass_solar_masses, accretion_rate_solar_masses_per_year, efficiency\n",
    "    )\n",
    "    eddington_luminosity_solar = eddington_luminosity(mass_solar_masses)\n",
    "\n",
    "    # Calculate additional derived properties\n",
    "\n",
    "    # Eddington ratio (L/L_edd)\n",
    "    eddington_ratio = (\n",
    "        accretion_luminosity_solar / eddington_luminosity_solar\n",
    "        if eddington_luminosity_solar > 0\n",
    "        else 0\n",
    "    )\n",
    "\n",
    "    # Calculate photon sphere radius (where photons orbit the black hole)\n",
    "    photon_sphere_radius_km = 1.5 * event_horizon_radius_km\n",
    "\n",
    "    # Calculate innermost stable circular orbit (ISCO)\n",
    "    innermost_stable_orbit_km = 3.0 * event_horizon_radius_km\n",
    "\n",
    "    # Estimate black hole evaporation time due to Hawking radiation\n",
    "    # Constants\n",
    "    import math\n",
    "\n",
    "    G = 6.67430e-11  # Gravitational constant in m^3 kg^-1 s^-2\n",
    "    c = 299792458  # Speed of light in m/s\n",
    "    hbar = 1.054571817e-34  # Reduced Planck constant in J·s\n",
    "    M_sun = 1.989e30  # Solar mass in kg\n",
    "\n",
    "    # Calculation (approximation based on Stephen Hawking's work)\n",
    "    # t_evap ≈ 5120 * π * G^2 * M^3 / (hbar * c^4)\n",
    "    mass_kg = mass_solar_masses * M_sun\n",
    "    evaporation_time_seconds = 5120 * math.pi * G**2 * mass_kg**3 / (hbar * c**4)\n",
    "    seconds_per_year = 365.25 * 24 * 3600\n",
    "    time_to_evaporation_years = evaporation_time_seconds / seconds_per_year\n",
    "\n",
    "    # Return the comprehensive analysis\n",
    "    return {\n",
    "        \"event_horizon_radius_km\": event_horizon_radius_km,\n",
    "        \"event_horizon_area_km2\": event_horizon_area_km2,\n",
    "        \"hawking_temperature_nK\": hawking_temperature_nK,\n",
    "        \"accretion_luminosity_solar\": accretion_luminosity_solar,\n",
    "        \"eddington_luminosity_solar\": eddington_luminosity_solar,\n",
    "        \"eddington_ratio\": eddington_ratio,\n",
    "        \"time_to_evaporation_years\": time_to_evaporation_years,\n",
    "        \"photon_sphere_radius_km\": photon_sphere_radius_km,\n",
    "        \"innermost_stable_orbit_km\": innermost_stable_orbit_km,\n",
    "    }\n",
    "\n",
    "\n",
    "# Register the wrapped function\n",
    "wrapped_function = client.create_wrapped_function(\n",
    "    primary_func=black_hole_analyzer,\n",
    "    functions=[\n",
    "        schwarzschild_radius,\n",
    "        event_horizon_area,\n",
    "        hawking_temperature,\n",
    "        accretion_luminosity,\n",
    "        eddington_luminosity,\n",
    "    ],\n",
    "    catalog=CATALOG,\n",
    "    schema=SCHEMA,\n",
    "    replace=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbc85fbc-b0ff-41fd-b94b-f89034631026",
   "metadata": {},
   "source": [
    "## Executing Wrapped Functions: Testing the Final Product\n",
    "\n",
    "The final stage in our Unity Catalog AI development workflow is executing and testing our wrapped function. This section demonstrates how to call our consolidated `black_hole_analyzer` function and interpret its comprehensive results.\n",
    "\n",
    "We'll analyze three different classes of black holes to showcase the power of our wrapped function:\n",
    "\n",
    "1. **Stellar Mass Black Hole (10 M☉)** - Similar to famous black holes like Cygnus X-1, these form from the collapse of massive stars\n",
    "2. **Intermediate Mass Black Hole (1,000 M☉)** - A rarer class that might form in dense star clusters\n",
    "3. **Supermassive Black Hole (4 million M☉)** - Like Sagittarius A* at the center of our Milky Way galaxy\n",
    "\n",
    "The execution is straightforward - we call the function with the appropriate parameters and receive a comprehensive analysis of each black hole's properties. Note how the wrapper function automatically leverages all the individual component functions while adding additional derived calculations.\n",
    "\n",
    "This example illustrates why the Unity Catalog AI 0.3.0 enhancements are so powerful for scientific computing workflows:\n",
    "\n",
    "1. **Modularity**: Each specialized calculation is maintained in its own function\n",
    "2. **Consolidation**: The wrapper brings everything together without code duplication\n",
    "3. **Extensibility**: New calculations can be added to either component functions or the wrapper\n",
    "4. **Clarity**: Complex physics is organized into logical units with clear documentation\n",
    "\n",
    "For data scientists and researchers, this pattern dramatically improves development efficiency while maintaining the specialized nature of scientific calculations. It's particularly valuable when creating functions for AI assistants, as it allows you to present complex scientific calculations through well-documented, logically structured interfaces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2570de9f-fc74-4f78-aca2-0fd29704246c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stellar Mass Black Hole (10 M☉):\n",
      "{'accretion_luminosity_solar': 14797.906972558185,\n",
      " 'eddington_luminosity_solar': 328605.84442117484,\n",
      " 'eddington_ratio': 0.045032391309485276,\n",
      " 'event_horizon_area_km2': 10966.500359645965,\n",
      " 'event_horizon_radius_km': 29.54126555055405,\n",
      " 'hawking_temperature_nK': 6.168429712630828,\n",
      " 'innermost_stable_orbit_km': 88.62379665166216,\n",
      " 'photon_sphere_radius_km': 44.31189832583108,\n",
      " 'time_to_evaporation_years': 2.0973585980140657e+70}\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\n",
      "Intermediate Mass Black Hole (1,000 M☉):\n",
      "{'accretion_luminosity_solar': 1479790.6972558186,\n",
      " 'eddington_luminosity_solar': 32860584.442117486,\n",
      " 'eddington_ratio': 0.045032391309485276,\n",
      " 'event_horizon_area_km2': 109665003.59645963,\n",
      " 'event_horizon_radius_km': 2954.126555055405,\n",
      " 'hawking_temperature_nK': 0.061684297126308275,\n",
      " 'innermost_stable_orbit_km': 8862.379665166214,\n",
      " 'photon_sphere_radius_km': 4431.189832583107,\n",
      " 'time_to_evaporation_years': 2.097358598014065e+76}\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\n",
      "Supermassive Black Hole (4 million M☉, like Sgr A*):\n",
      "{'accretion_luminosity_solar': 0.0,\n",
      " 'eddington_luminosity_solar': 131442337768.46992,\n",
      " 'eddington_ratio': 0.0,\n",
      " 'event_horizon_area_km2': 1754640057543353.5,\n",
      " 'event_horizon_radius_km': 11816506.220221618,\n",
      " 'hawking_temperature_nK': 1.5421074281577068e-05,\n",
      " 'innermost_stable_orbit_km': 35449518.66066486,\n",
      " 'photon_sphere_radius_km': 17724759.33033243,\n",
      " 'time_to_evaporation_years': 1.342309502729002e+87}\n"
     ]
    }
   ],
   "source": [
    "# Execute our wrapped black hole analyzer function with a variety of black hole masses\n",
    "import ast\n",
    "from pprint import pprint\n",
    "\n",
    "# Stellar mass black hole (10 solar masses, like Cygnus X-1)\n",
    "stellar_bh_result = client.execute_function(\n",
    "    function_name=f\"{CATALOG}.{SCHEMA}.black_hole_analyzer\",\n",
    "    parameters={\n",
    "        \"mass_solar_masses\": 10.0,\n",
    "        \"accretion_rate_solar_masses_per_year\": 1e-8,\n",
    "        \"efficiency\": 0.1,\n",
    "    },\n",
    ")\n",
    "print(\"Stellar Mass Black Hole (10 M☉):\")\n",
    "pprint(ast.literal_eval(stellar_bh_result.value))\n",
    "print(\"\\n\" + \"-\" * 80 + \"\\n\")\n",
    "\n",
    "# Intermediate mass black hole (1,000 solar masses)\n",
    "intermediate_bh_result = client.execute_function(\n",
    "    function_name=f\"{CATALOG}.{SCHEMA}.black_hole_analyzer\",\n",
    "    parameters={\"mass_solar_masses\": 1000.0, \"accretion_rate_solar_masses_per_year\": 1e-6},\n",
    ")\n",
    "print(\"Intermediate Mass Black Hole (1,000 M☉):\")\n",
    "pprint(ast.literal_eval(intermediate_bh_result.value))\n",
    "print(\"\\n\" + \"-\" * 80 + \"\\n\")\n",
    "\n",
    "# Supermassive black hole (4 million solar masses, like Sagittarius A*)\n",
    "smbh_result = client.execute_function(\n",
    "    function_name=f\"{CATALOG}.{SCHEMA}.black_hole_analyzer\", parameters={\"mass_solar_masses\": 4.0e6}\n",
    ")\n",
    "print(\"Supermassive Black Hole (4 million M☉, like Sgr A*):\")\n",
    "pprint(ast.literal_eval(smbh_result.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8dff208-11bc-495e-9797-7bbd74eeac05",
   "metadata": {},
   "source": [
    "## Sandbox Protection: Safeguarding Against Dangerous Executions\n",
    "\n",
    "One of the most powerful security enhancements in Unity Catalog AI 0.3.0 is the improved sandbox execution mode. This feature provides critical protection for production systems by isolating potentially hazardous function executions in separate processes.\n",
    "\n",
    "### Why Sandbox Protection Matters\n",
    "\n",
    "When working with AI agents that can call arbitrary functions, you face a fundamental challenge: users might unknowingly (or deliberately) provide parameters that could:\n",
    "\n",
    "1. **Crash your system** with unbounded computational demands\n",
    "2. **Lock up resources** with infinite loops or excessive memory usage\n",
    "3. **Access sensitive system resources** through imported modules\n",
    "4. **Block other operations** by monopolizing computational resources\n",
    "\n",
    "The sandbox execution mode creates a protective barrier by running functions in isolated processes with strict resource limitations.\n",
    "\n",
    "### How Sandbox Protection Works\n",
    "\n",
    "In this section, we'll demonstrate how the sandbox works by creating a deliberately resource-intensive function: a Mandelbrot set calculator. While this function is perfectly safe with reasonable parameters, it can consume enormous computational resources when given large dimensions or iteration counts.\n",
    "\n",
    "The sandbox protection system:\n",
    "\n",
    "1. **Runs code in isolated processes** rather than in the main application process\n",
    "2. **Enforces CPU time limits** to prevent excessive computation\n",
    "3. **Restricts memory usage** to prevent memory exhaustion\n",
    "4. **Blocks access to sensitive system modules** like `os`, `sys`, and `subprocess`\n",
    "5. **Imposes execution timeouts** to ensure functions terminate\n",
    "\n",
    "For our demonstration, we've customized the timeout to just 3 seconds (compared to the default 30 seconds) to show how quickly the system can respond to potentially problematic executions.\n",
    "\n",
    "This protection is especially important for GenAI applications, where large language models might explore edge cases or extreme parameters that human users would typically avoid. By configuring the sandbox mode appropriately, you can ensure that even the most computationally demanding function calls will fail gracefully rather than bringing down your entire application."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e278f9b",
   "metadata": {},
   "source": [
    "## ⚠️ Computational Complexity Warning: Mandelbrot Set Calculator\n",
    "\n",
    "The `mandelbrot_set_calculator` function provides a fascinating mathematical tool for generating the famous Mandelbrot set fractal. However, it presents significant computational challenges that make it a perfect example of why unrestricted function execution in AI agents can be dangerous.\n",
    "\n",
    "### The Computational Complexity Problem\n",
    "\n",
    "#### Time Complexity\n",
    "- **O(width × height × max_iterations)** - This is a cubic growth pattern in the worst case\n",
    "- With maximum allowed parameters (5000×5000 pixels, 100,000 iterations), this results in **2.5 trillion operations**\n",
    "- A single call with large parameters can monopolize CPU resources for minutes or even hours\n",
    "\n",
    "#### Memory Consumption\n",
    "- Creates a NumPy array of size `width × height × 4 bytes` (for int32)\n",
    "- At maximum dimensions (5000×5000), requires **~95MB** just for the result matrix\n",
    "- Additional memory overhead for the Python lists created for JSON serialization\n",
    "\n",
    "### Dangers in GenAI Agent Contexts\n",
    "\n",
    "If an AI agent has unrestricted access to execute this function:\n",
    "\n",
    "1. **Denial of Service (DoS)**: A user could easily (intentionally or unintentionally) trigger a computation that paralyzes the main application process\n",
    "\n",
    "2. **Resource Starvation**: A single request with parameters like `width=5000, height=5000, max_iterations=100000` could consume all available CPU resources\n",
    "\n",
    "3. **Unpredictable Response Times**: Even with moderate parameters, execution time varies dramatically based on the specific region of the complex plane being calculated\n",
    "\n",
    "4. **Memory Exhaustion**: Large result matrices can consume significant memory, potentially causing OOM (Out of Memory) errors\n",
    "\n",
    "## Safeguard Solutions\n",
    "\n",
    "This is why process/queue-based sandbox execution with the following protections is essential:\n",
    "\n",
    "- **CPU Limits**: Restrict the maximum CPU time any single calculation can consume\n",
    "- **Memory Caps**: Prevent memory usage from exceeding predefined thresholds\n",
    "- **Execution Timeouts**: Automatically terminate calculations that exceed time limits\n",
    "- **Process Isolation**: Run computationally intensive functions in separate processes to protect the main application\n",
    "- **Queue Management**: Prevent queue flooding by limiting concurrent intensive operations\n",
    "\n",
    "### Practical Example\n",
    "\n",
    "Consider these two parameter sets:\n",
    "\n",
    "1. `width=100, height=100, max_iterations=1000`\n",
    "   - ~10 million operations\n",
    "   - Completes almost instantly\n",
    "   - Memory usage: ~40KB\n",
    "\n",
    "2. `width=5000, height=5000, max_iterations=50000`\n",
    "   - ~1.25 trillion operations\n",
    "   - Could take hours on standard hardware\n",
    "   - Memory usage: ~95MB\n",
    "\n",
    "Without proper safeguards, the second example could effectively crash or freeze your application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5a0fe60e-7f71-41ca-be1b-27786f3a04d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mandelbrot_set_calculator(\n",
    "    width: int,\n",
    "    height: int,\n",
    "    max_iterations: int = 1000,\n",
    "    x_min: float = -2.0,\n",
    "    x_max: float = 1.0,\n",
    "    y_min: float = -1.5,\n",
    "    y_max: float = 1.5,\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Calculates the Mandelbrot set for the specified parameters.\n",
    "\n",
    "    This function creates a 2D array representing the Mandelbrot set. Each element\n",
    "    in the array contains the number of iterations required to determine if the\n",
    "    corresponding point in the complex plane is in the Mandelbrot set.\n",
    "\n",
    "    WARNING: This function can be computationally expensive with large dimensions or\n",
    "    high iteration counts. Values above 2000x2000 pixels or 10000 iterations may\n",
    "    cause excessive CPU usage or memory consumption.\n",
    "\n",
    "    Args:\n",
    "        width (int): Width of the resulting image in pixels.\n",
    "            Valid range: 1 to 5000.\n",
    "        height (int): Height of the resulting image in pixels.\n",
    "            Valid range: 1 to 5000.\n",
    "        max_iterations (int, optional): Maximum number of iterations to determine if a point\n",
    "            is in the set. Defaults to 1000.\n",
    "            Valid range: 10 to 100000.\n",
    "        x_min (float, optional): Minimum real value of the complex plane. Defaults to -2.0.\n",
    "        x_max (float, optional): Maximum real value of the complex plane. Defaults to 1.0.\n",
    "        y_min (float, optional): Minimum imaginary value of the complex plane. Defaults to -1.5.\n",
    "        y_max (float, optional): Maximum imaginary value of the complex plane. Defaults to 1.5.\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary containing:\n",
    "            - dimensions: Tuple of (width, height)\n",
    "            - iterations_matrix: 2D list of iteration counts\n",
    "            - calculation_info: Information about the calculation parameters\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If width, height, or max_iterations are outside their valid ranges.\n",
    "\n",
    "    Example:\n",
    "        >>> result = mandelbrot_set_calculator(100, 100, 500)\n",
    "        >>> print(\n",
    "        ...     f\"Generated a Mandelbrot set of {result['dimensions'][0]}x{result['dimensions'][1]} pixels\"\n",
    "        ... )\n",
    "    \"\"\"\n",
    "    # Input validation\n",
    "    if not (1 <= width <= 5000):\n",
    "        raise ValueError(\"Width must be between 1 and 5000 pixels\")\n",
    "    if not (1 <= height <= 5000):\n",
    "        raise ValueError(\"Height must be between 1 and 5000 pixels\")\n",
    "    if not (10 <= max_iterations <= 100000):\n",
    "        raise ValueError(\"Max iterations must be between 10 and 100000\")\n",
    "\n",
    "    # Import NumPy here to keep it isolated to this function\n",
    "    import numpy as np\n",
    "\n",
    "    # Create a 2D array to store the iteration counts\n",
    "    result_matrix = np.zeros((height, width), dtype=np.int32)\n",
    "\n",
    "    # Calculate the step size in the complex plane\n",
    "    x_step = (x_max - x_min) / width\n",
    "    y_step = (y_max - y_min) / height\n",
    "\n",
    "    # Loop through each pixel\n",
    "    for y in range(height):\n",
    "        for x in range(width):\n",
    "            # Convert pixel coordinate to complex number\n",
    "            c_real = x_min + x * x_step\n",
    "            c_imag = y_min + y * y_step\n",
    "            c = complex(c_real, c_imag)\n",
    "\n",
    "            # Perform Mandelbrot iteration\n",
    "            z = complex(0, 0)\n",
    "            iteration = 0\n",
    "\n",
    "            while abs(z) <= 2 and iteration < max_iterations:\n",
    "                z = z**2 + c\n",
    "                iteration += 1\n",
    "\n",
    "            # Store the iteration count\n",
    "            result_matrix[y, x] = iteration\n",
    "\n",
    "    # Convert the NumPy array to a Python list for JSON serialization\n",
    "    iterations_list = result_matrix.tolist()\n",
    "\n",
    "    # Create the return dictionary\n",
    "    result = {\n",
    "        \"dimensions\": (width, height),\n",
    "        \"iterations_matrix\": iterations_list,\n",
    "        \"calculation_info\": {\n",
    "            \"max_iterations\": max_iterations,\n",
    "            \"complex_plane_bounds\": {\n",
    "                \"x_min\": x_min,\n",
    "                \"x_max\": x_max,\n",
    "                \"y_min\": y_min,\n",
    "                \"y_max\": y_max,\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7974ee96-0498-484b-9a31-194c3cc27847",
   "metadata": {},
   "source": [
    "## Sandbox Protection: Safeguarding Against Runaway Computations\n",
    "\n",
    "One of the most significant security enhancements in Unity Catalog AI 0.3.0 is the improved sandbox execution mode. This feature provides critical protection for production systems when integrating with AI agents or handling user-provided parameters.\n",
    "\n",
    "### The Problem: Dangerous Function Executions\n",
    "\n",
    "When exposing functions to AI agents or end users, several risks emerge:\n",
    "\n",
    "1. **Resource Exhaustion**: Functions with unbounded computation can consume all available CPU or memory\n",
    "2. **Denial of Service**: Poorly parameterized functions can effectively lock up your system\n",
    "3. **System Access**: Malicious code might attempt to access sensitive system resources\n",
    "4. **Infinite Loops**: Functions without proper termination conditions can run indefinitely\n",
    "\n",
    "These issues are particularly concerning with GenAI applications because large language models may explore edge cases or suggest extreme parameters that could unintentionally stress your system.\n",
    "\n",
    "### The Solution: Process Isolation with Resource Limits\n",
    "\n",
    "The sandbox execution mode in Unity Catalog AI 0.3.0 provides comprehensive protection through several mechanisms:\n",
    "\n",
    "1. **Process Isolation**: Functions run in separate processes rather than in your main application\n",
    "2. **Resource Limitations**:\n",
    "  - **CPU Time Limits**: Restricts total CPU time consumption\n",
    "  - **Memory Caps**: Prevents memory exhaustion\n",
    "  - **Execution Timeouts**: Forces termination after a specified wall clock time\n",
    "3. **Module Restrictions**: Blocks access to dangerous system modules like `os`, `sys`, and `subprocess`\n",
    "4. **File Access Prevention**: Disables the built-in `open()` function to prevent filesystem access\n",
    "\n",
    "### Configuration Through Environment Variables\n",
    "\n",
    "The sandbox protection system is highly configurable through environment variables:\n",
    "\n",
    "- `EXECUTOR_TIMEOUT`: Maximum wall clock time (in seconds) before termination\n",
    "- `EXECUTOR_MAX_CPU_TIME_LIMIT`: Maximum CPU time (in seconds)\n",
    "- `EXECUTOR_MAX_MEMORY_LIMIT`: Maximum memory usage (in MB)\n",
    "- `EXECUTOR_DISALLOWED_MODULES`: Comma-separated list of modules to block\n",
    "\n",
    "As demonstrated in our example, you can adjust these settings to match your specific requirements and risk tolerance.\n",
    "\n",
    "### The Mandelbrot Example\n",
    "\n",
    "Our Mandelbrot set calculator provides a perfect demonstration of how sandbox protection works:\n",
    "\n",
    "- With modest parameters (50x50 pixels), the function completes quickly and returns valid results\n",
    "- With extreme parameters (3000x3000 pixels), the computation would normally consume significant resources and time\n",
    "- The sandbox detects the excessive computation and terminates it after our configured 3-second timeout\n",
    "\n",
    "This protection mechanism ensures that even if an AI agent or user submits problematic parameters, your system remains responsive and stable. The function simply fails gracefully with a timeout message rather than bringing down your entire application.\n",
    "\n",
    "For production GenAI applications, this layer of protection is essential when exposing computational functions to large language models, which might not inherently understand the resource implications of the parameters they suggest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d293c720-b828-45ff-8474-98ce96bd2318",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Registered function: astro_catalog.physics_functions.mandelbrot_set_calculator\n",
      "Original execution timeout: 30 seconds\n",
      "Modified execution timeout: 3 seconds\n",
      "\n",
      "--- Safe Parameters Test ---\n",
      "Safe execution successful! Generated a 50x50 matrix\n",
      "\n",
      "--- Dangerous Parameters Test ---\n",
      "Attempting execution with parameters that will exceed the timeout...\n",
      "Result: The function execution has timed out and has been canceled due to excessive resource consumption.\n",
      "There are two timeout conditions to consider:\n",
      "\t1. You can increase the CPU execution timeout by setting the environment variable EXECUTOR_MAX_CPU_TIME_LIMIT. The default value is 10 seconds.\n",
      "\t2. You can increase the wall-clock time limit by setting the environment variable EXECUTOR_TIMEOUT. The default value is 20 seconds.\n",
      "\n",
      "\n",
      "Restored execution timeout to original value: 30 seconds\n"
     ]
    }
   ],
   "source": [
    "# Register the Mandelbrot function\n",
    "client.create_python_function(\n",
    "    catalog=CATALOG, schema=SCHEMA, func=mandelbrot_set_calculator, replace=True\n",
    ")\n",
    "\n",
    "print(f\"Registered function: {CATALOG}.{SCHEMA}.mandelbrot_set_calculator\")\n",
    "\n",
    "# Import the required modules for environment configuration\n",
    "import os\n",
    "\n",
    "# Save the original timeout value if it exists\n",
    "original_timeout = os.environ.get(\"EXECUTOR_TIMEOUT\", \"30\")\n",
    "print(f\"Original execution timeout: {original_timeout} seconds\")\n",
    "\n",
    "# Modify the timeout to a lower value for demonstration\n",
    "# This would terminate computations after just 3 seconds\n",
    "os.environ[\"EXECUTOR_TIMEOUT\"] = \"3\"\n",
    "print(f\"Modified execution timeout: {os.environ['EXECUTOR_TIMEOUT']} seconds\")\n",
    "\n",
    "try:\n",
    "    print(\"\\n--- Safe Parameters Test ---\")\n",
    "    # Execute with safe parameters\n",
    "    safe_result = client.execute_function(\n",
    "        function_name=f\"{CATALOG}.{SCHEMA}.mandelbrot_set_calculator\",\n",
    "        parameters={\"width\": 50, \"height\": 50, \"max_iterations\": 100},\n",
    "    )\n",
    "    import ast\n",
    "\n",
    "    safe_result_dict = ast.literal_eval(safe_result.value)\n",
    "    print(\n",
    "        f\"Safe execution successful! Generated a {safe_result_dict['dimensions'][0]}x{safe_result_dict['dimensions'][1]} matrix\"\n",
    "    )\n",
    "\n",
    "    print(\"\\n--- Dangerous Parameters Test ---\")\n",
    "    # Execute with parameters that would timeout\n",
    "    print(\"Attempting execution with parameters that will exceed the timeout...\")\n",
    "    dangerous_result = client.execute_function(\n",
    "        function_name=f\"{CATALOG}.{SCHEMA}.mandelbrot_set_calculator\",\n",
    "        parameters={\"width\": 3000, \"height\": 3000, \"max_iterations\": 5000},\n",
    "    )\n",
    "    print(f\"Result: {dangerous_result.error}\")\n",
    "\n",
    "except Exception as e:\n",
    "    print(f\"An error occurred: {str(e)}\")\n",
    "finally:\n",
    "    # Restore the original timeout\n",
    "    os.environ[\"EXECUTOR_TIMEOUT\"] = original_timeout\n",
    "    print(\n",
    "        f\"\\nRestored execution timeout to original value: {os.environ['EXECUTOR_TIMEOUT']} seconds\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff3defa1-56b2-409f-9c07-f6952477a54d",
   "metadata": {},
   "source": [
    "## Conclusion: Building Robust Scientific Functions with Unity Catalog AI\n",
    "\n",
    "In this tutorial, we've explored the powerful development workflow enhancements introduced in Unity Catalog AI versions 0.2.0 and 0.3.0. These features significantly improve the development experience when creating, testing, and deploying scientific functions for AI-powered applications.\n",
    "\n",
    "### Key Workflow Improvements We've Covered\n",
    "\n",
    "1. **Enhanced Function Management**\n",
    "  - Setting up connections to Unity Catalog\n",
    "  - Creating and registering complex scientific functions\n",
    "  - Retrieving function source code for inspection and modification\n",
    "  - Updating functions with new capabilities\n",
    "\n",
    "2. **Function Consolidation**\n",
    "  - Creating independent scientific calculation functions\n",
    "  - Building wrapper functions that consolidate multiple calculations\n",
    "  - Minimizing code duplication through function wrapping\n",
    "  - Providing comprehensive interfaces for AI agents\n",
    "\n",
    "3. **Safety and Security**\n",
    "  - Executing functions in sandbox mode for process isolation\n",
    "  - Configuring resource limits to prevent runaway computations\n",
    "  - Protecting against module imports that could access system resources\n",
    "  - Ensuring graceful failure for problematic parameter sets\n",
    "\n",
    "### Practical Applications\n",
    "\n",
    "The techniques demonstrated in this tutorial are particularly valuable for:\n",
    "\n",
    "- **Scientific Research**: Building modular, reusable calculation libraries\n",
    "- **Data Science Workflows**: Creating robust parameter exploration functions\n",
    "- **AI Agent Integration**: Providing specialized domain knowledge to LLMs\n",
    "- **Educational Tools**: Developing interactive scientific demonstrations\n",
    "\n",
    "### Next Steps\n",
    "\n",
    "To further explore Unity Catalog AI's capabilities:\n",
    "\n",
    "1. **Explore Integration Packages**: Try connecting your functions to AI frameworks like LangChain, LlamaIndex, or directly to models from OpenAI, Anthropic, or Gemini\n",
    "2. **Build Domain-Specific Libraries**: Create your own specialized function collections for your field\n",
    "3. **Implement Function Versioning**: Use the catalog and schema structure to maintain different versions of your functions\n",
    "4. **Explore Production Deployment**: Move from local testing to serverless production deployment\n",
    "\n",
    "By leveraging these powerful development workflow enhancements, you can create more maintainable, secure, and effective scientific functions that seamlessly integrate with modern AI systems."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
